# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

using ForwardDiff
using Printf
using LinearAlgebra

const Vec = Vector{Float64}
const Mat = Matrix{Float64}
const iMat = Matrix{Int}

N_CAM_PARAMS = 11
ROT_IDX = 1
C_IDX = 4
F_IDX = 7
X0_IDX = 8
RAD_IDX = 10

##################### objective #############################

function rodrigues_rotate_point(rot :: Vector{T}, X :: Vector{T}) where T
  sqtheta = sum(rot.*rot)
  if sqtheta > 1e-10
    theta = sqrt(sqtheta)
    costheta = cos(theta)
    sintheta = sin(theta)
    theta_inverse = 1. / theta

    w = theta_inverse * rot
    w_cross_X = cross(w,X)
    tmp = dot(w,X) * (1. - costheta)

    X*costheta + w_cross_X * sintheta + w * tmp
  else
    X + cross(rot,X)
  end
end

function radial_distort(rad_params,proj)
  rsq = sum(proj.*proj)
  L = 1. + rad_params[1]*rsq + rad_params[2]*rsq*rsq
  proj*L
end

function project(cam,X)
  Xcam = rodrigues_rotate_point(cam[ROT_IDX:ROT_IDX+2], X - cam[C_IDX:C_IDX+2])
  distorted = radial_distort(cam[RAD_IDX:RAD_IDX+1], Xcam[1:2]/Xcam[3])
  distorted*cam[F_IDX] + cam[X0_IDX:X0_IDX+1]
end

function compute_reproj_err(cam,X,w,feat)
  return w*(project(cam,X) - feat)
end

function ba_objective(cams,X,w,obs,feats)
  reproj_err = similar(feats)
  for i in 1:size(feats,2)
    reproj_err[:,i] = compute_reproj_err(cams[:,obs[1,i]],X[:,obs[2,i]],w[i],feats[:,i])
  end
  w_err = 1.0 .- w.*w
  (reproj_err, w_err)
end

#################### derivatives extra ##########################

function pack(cam,X,w)
  [cam[:];X[:];w]
end

function unpack(packed)
  packed[1:end-4],packed[end-3:end-1],packed[end]
end

function compute_w_err(w)
    1.0 - w*w
end
compute_w_err_d = x -> ForwardDiff.derivative(compute_w_err, x)

function compute_reproj_err_d(params, feat)
  cam, X, w = unpack(params)
  compute_reproj_err(cam,X,w,feat)
end

function compute_ba_J(cams,X,w,obs,feats)
  p = size(obs,2)
  reproj_err_d = zeros(2*p, N_CAM_PARAMS + 3 + 1)
  for i in 1:p
    compute_reproj_err_d_i = x -> compute_reproj_err_d(x, feats[:,i])
    idx = (2*(i-1))+1
    reproj_err_d[idx:idx+1,:] = ForwardDiff.jacobian(compute_reproj_err_d_i,
              pack(cams[:,obs[1,i]],X[:,obs[2,i]],w[i]))
  end
  w_err_d = zeros(1,p)
  for i in 1:p
    w_err_d[i] = compute_w_err_d(w[i])
  end
  (reproj_err_d, w_err_d)
end

# precompile as much as possible
precompile(rodrigues_rotate_point,(Vec,Vec))
precompile(radial_distort,(Vec,Vec))
precompile(project,(Vec,Vec))
precompile(compute_reproj_err,(Vec,Vec,Float64,Vec))
precompile(ba_objective,(Mat,Mat,Vec,iMat,Mat))

precompile(compute_ba_J,(Mat,Mat,Vec,iMat,Mat))
precompile(compute_reproj_err_d,(Vec,Vec))
precompile(compute_w_err_d,(Vec,))

#################### IO ############################
function read_ba_instance(fn)
  fid = open(fn)
  lines = readlines(fid)
  close(fid)
  line=split(lines[1]," ")
  n = parse(Int,line[1])
  m = parse(Int,line[2])
  p = parse(Int,line[3])
  off = 2

  one_cam = zeros(Float64,N_CAM_PARAMS,1)
  line=split(lines[off]," ")
  for i in 1:N_CAM_PARAMS
    one_cam[i] = parse(Float64,line[i])
  end
  cams = repeat(one_cam,1,n)
  off += 1

  one_X = zeros(Float64,3,1)
  line=split(lines[off]," ")
  for i in 1:3
    one_X[i] = parse(Float64,line[i])
  end
  X = repeat(one_X,1,m)
  off += 1

  one_w = parse(Float64,lines[off])
  w = repeat([one_w],1,p)
  off += 1

  one_feat = zeros(Float64,2,1)
  line=split(lines[off]," ")
  for i in 1:2
    one_feat[i] = parse(Float64,line[i])
  end
  feats = repeat(one_feat,1,p)

  camIdx = 1
  ptIdx = 1
  obs = zeros(Int,2,p)
  for i in 1:p
    obs[1,i] = camIdx
    obs[2,i] = ptIdx
    camIdx = (camIdx%n) + 1
    ptIdx = (ptIdx%m) + 1
  end

  (cams,X,w,obs,feats)
end

function write_times(fn,tf,tJ)
  @printf "Writing to %s\n" fn
  fid = open(fn,"w")
  @printf fid "%f %f\r\n" tf tJ
  @printf fid "tf tJ\r\n"
  close(fid)
end

##################### run it ###########################

# Read instance
function main(ARGS)
  dir_in = ARGS[1]
  dir_out = ARGS[2]
  fn = ARGS[3]
  nruns_f = parse(Int,ARGS[4])
  nruns_J = parse(Int,ARGS[5])

  fn_in = string(dir_in, fn)
  fn_out = string(dir_out, fn)

  cams,X,w,obs,feats = read_ba_instance(string(fn_in,".txt"))

  # compute once for precompilation
  reproj_err, w_err = ba_objective(cams,X,w,obs,feats)
  # Time runs
  tf = @elapsed for i in 1:nruns_f
    reproj_err, w_err = ba_objective(cams,X,w,obs,feats)
  end
  tf /= nruns_f
  @printf "tf: %g\n" tf
  #println(reproj_err)
  #println(w_err)

  # compute once for precompilation
  J = compute_ba_J(cams,X,w,obs,feats)
  # Time runs
  tJ = @elapsed for i in 1:nruns_J
    J = compute_ba_J(cams,X,w,obs,feats)
  end
  tJ /= nruns_J;
  @printf "tJ: %g\n" tJ
  #println("J:")
  #println(J)

  name = "Julia"

  #write_J(string(fn_out,"_J_",name,".txt"),J)
  write_times(string(fn_out,"_times_",name,".txt"),tf,tJ)
end

if isinteractive()
  main([raw"C:\dev\github\autodiff/data/ba/", raw"C:\dev\github\autodiff/tmp/Release/ba/Julia/", "ba1", "1", "1", "60"])
else
  main(ARGS)
end


export minreal

"""
### function minreal(G::TransferFunction)

Compute the minimal realization of a dynamic system represented by a transfer
function.  In other words, compute a transfer function without common poles and
zeros.

##### Args

* G: Transfer function.

##### Returns

The transfer function of the minimal realization of the dynamic system.

"""

function minreal(G::TransferFunction)
    # Get the list of poles and zeros.
    p = pole(G)
    z = izero(G)

    # Get the system gain.
    k = G.num[end]/G.den[end]

    # Get the poles and zeros of the minimal system.
    mp = p
    mz = Array{Complex128}(0)

    # Auxiliary variable to compute tolarance.
    #
    # This algorithm was adapted from file `__minreal__.m` of Octave Forge.
    sqrt_eps = sqrt(eps())

    # Remove the poles that can be cancelled with zeros.
    for i = 1:length(z)
        # Compute the tolerance.
        tol = (abs(z[i]) < sqrt_eps) ? 10000*sqrt_eps : 10000*abs(z[i])*sqrt_eps

        # Compute the distance from the poles to the zero.
        Δ = abs.(p-z[i])

        # Find if there are any poles within a distance of `tol` from the
        # analyzed zero.
        ind = find(Δ .< tol)

        if !isempty(ind)
            # Delete the closest pole to the current zero.
            deleteat!(mp, indmin(Δ))
        else
            push!(mz,round(z[i],6))
        end
    end

    # Return the minimal transfer function.
    zpk(mz, mp, k)
end


"Fresnel reflection, needs pol_type soon too!"
function Fᵣ(n::FT,θᵢ::FT) where FT
    #@show rad2deg(θᵢ)
    # Refractive index of air
    nᵢ = FT(1)
    # Angle of transmitted light in medium:
    θₜ = asin(nᵢ * sin(θᵢ)/n)
    #(sin(α - θₛ)^2 / sin(α + θₛ)^2 +  tan(α - θₛ)^2 / tan(α + θₛ)^2)/2
    r_perpendicular = (nᵢ * cos(θᵢ) - n * cos(θₜ)) / (nᵢ * cos(θᵢ) + n * cos(θₜ))
    r_parallel      = (nᵢ * cos(θₜ) - n * cos(θᵢ)) / (nᵢ * cos(θₜ) + n * cos(θᵢ))
    return (r_perpendicular^2 +  r_parallel^2)/2
end


export Likelihood

@doc raw"""
    Likelihood(M<:ParameterizedMeasure, x)

"Observe" a value `x`, yielding a function from the parameters to ℝ.

Likelihoods are most commonly used in conjunction with an existing _prior_
measure to yield a new measure, the _posterior_. In Bayes's Law, we have

``P(θ|x) ∝ P(θ) P(x|θ)``

Here ``P(θ)`` is the prior. If we consider ``P(x|θ)`` as a function on ``θ``,
then it is called a likelihood.

Since measures are most commonly manipulated using `density` and `logdensity`,
it's awkward to commit a (log-)likelihood to using one or the other. To evaluate
a `Likelihood`, we therefore use `density` or `logdensity`, depending on the
circumstances. In the latter case, it is of course acting as a log-density.

For example,

    julia> ℓ = Likelihood(Normal{(:μ,)}, 2.0)
    Likelihood(Normal{(:μ,), T} where T, 2.0)

    julia> density(ℓ, (μ=2.0,))
    1.0

    julia> logdensity(ℓ, (μ=2.0,))
    -0.0

If, as above, the measure includes the parameter information, we can optionally
leave it out of the second argument in the call to `density` or `logdensity`.

    julia> density(ℓ, 2.0)
    1.0

    julia> logdensity(ℓ, 2.0)
    -0.0

With several parameters, things work as expected:

    julia> ℓ = Likelihood(Normal{(:μ,:σ)}, 2.0)
    Likelihood(Normal{(:μ, :σ), T} where T, 2.0)

    julia> logdensity(ℓ, (μ=2, σ=3))
    -1.0986122886681098

    julia> logdensity(ℓ, (2,3))
    -1.0986122886681098

    julia> logdensity(ℓ, [2, 3])
    -1.0986122886681098

---------

    Likelihood(M<:ParameterizedMeasure, constraint::NamedTuple, x)

In some cases the measure might have several parameters, and we may want the
(log-)likelihood with respect to some subset of them. In this case, we can use
the three-argument form, where the second argument is a constraint. For example,

    julia> ℓ = Likelihood(Normal{(:μ,:σ)}, (σ=3.0,), 2.0)
    Likelihood(Normal{(:μ, :σ), T} where T, (σ = 3.0,), 2.0)

Similarly to the above, we have

    julia> density(ℓ, (μ=2.0,))
    0.3333333333333333

    julia> logdensity(ℓ, (μ=2.0,))
    -1.0986122886681098

    julia> density(ℓ, 2.0)
    0.3333333333333333

    julia> logdensity(ℓ, 2.0)
    -1.0986122886681098

-----------------------

Finally, let's return to the expression for Bayes's Law,

``P(θ|x) ∝ P(θ) P(x|θ)``

The product on the right side is computed pointwise. To work with this in
MeasureTheory, we have a "pointwise product" `⊙`, which takes a measure and a
likelihood, and returns a new measure, that is, the unnormalized posterior that has density ``P(θ) P(x|θ)`` with respect to the base measure of the prior.

For example, say we have

    μ ~ Normal()
    x ~ Normal(μ,σ)
    σ = 1

and we observe `x=3`. We can compute the posterior measure on `μ` as

    julia> post = Normal() ⊙ Likelihood(Normal{(:μ, :σ)}, (σ=1,), 3)
    Normal() ⊙ Likelihood(Normal{(:μ, :σ), T} where T, (σ = 1,), 3)

    julia> logdensity(post, 2)
    -2.5
"""
struct Likelihood{F,X}
    f::F
    x::X
end

function Base.show(io::IO, ℓ::Likelihood{Tuple{M,NamedTuple{N,T}}}) where {M<:Type, N,T}
    (m,c) = ℓ.f
    println(io, "Likelihood(",m,", ",c,", ", ℓ.x, ")")
end

function Base.show(io::IO, ℓ::Likelihood)
    println(io, "Likelihood(",ℓ.f, ", ", ℓ.x, ")")
end

Likelihood(μ::M, x) where {M<:AbstractMeasure} = Likelihood(M, x)

function Likelihood(::Type{M}, constraint::NamedTuple, x) where {M <: ParameterizedMeasure}
    Likelihood((M, constraint), x)
end

function Likelihood(μ::M, constraint::NamedTuple, x) where {M<:AbstractMeasure}
    Likelihood((M, constraint), x)
end

function logdensity(ℓ::Likelihood{Tuple{M,NT}}, p::NamedTuple) where {M<:Type, NT <: NamedTuple}
    (D, constraint) = ℓ.f
    return logdensity(D(merge(p, constraint)), ℓ.x)
end

function logdensity(ℓ::Likelihood{Tuple{M,NT}}, p) where {M<:Type, NT <: NamedTuple}
    freevar = params(ℓ.f...)
    logdensity(ℓ, NamedTuple{freevar}(p))
end

logdensity(ℓ::Likelihood, p) = logdensity(ℓ.f(p), ℓ.x)


# # Introduction

# We first load [DrWatson](https://juliadynamics.github.io/DrWatson.jl/dev/),
# a scientific project management software that helps reproducibility.
using DrWatson

# We then activate this project, which also loads a bundle of packages
# (most notably ClimateBase.jl), and includes all files in the `src` folder.
@quickactivate :AlbedoProperties

# We then load the style file for the figures (text size, colors, etc.)
include(srcdir("plotting", "advanced.jl"))
include(papersdir("paperstyle.jl"))

# Load all necessary CERES files explicitly
EBAF_TOA = datadir("CERES", "EBAF_TOA.nc")
EBAF_SFC = datadir("CERES", "EBAF_SFC.nc")
SEAICE = datadir("CERES", "SYN1deg_AUX.nc")

# And from these files load all spatiotemporal fields necessary
R = ncread(EBAF_TOA, "toa_sw_all_mon")
K = ncread(EBAF_TOA, "toa_sw_clr_c_mon")
I = ncread(EBAF_TOA, "solar_mon")
F = ncread(EBAF_TOA, "cldarea_total_daynight_mon") ./ 100
T = ncread(EBAF_TOA, "cldtau_total_day_mon")
O = ncread(SEAICE, "aux_ocean_mon") ./ 100
E = ncread(SEAICE, "aux_snow_mon") ./ 100
L = clamp.(1 .- O .- E, 0, 1) # ice-free land
R

# The time vector is
t = dims(R, Time).val

# and the month counter is
tn = 1:length(t)

# And the following index is the most amount of data we can use for full-years
mys = maxyearspan(t)

# Notice that not all CERES data have this time span. Surface fluxes and ice fractions
# cover less amount of time, thus for operations involving this data we will infact be
# forced to use a smaller time span.

# Time in cumulative number of days per month is
truetime = daysinmonth.(t)
truetime = Float32.(cumsum(truetime))

# Ιn most figures we will be showing a sub-set of the total time (purely for visual clarity,
# and actual span chosen below is completely arbitrary)
tspan = (Date(2008), Date(2015))

# This is how we decompose
# some fields into northern-southern hemisphere average timeseries.
Ιn, Is = hemispheric_means(I)
Rn, Rs = hemispheric_means(R)
Rn

# This is a good point to demonstrate how critical it is for one to do proper temporal
# averaging. Integrated over a full yearly cycle, both hemispheres receive
# the same amount of insolation. This is a direct consequence of Kepler's second law.
timemean(Ιn) - timemean(Is)

# Using a standard mean instead, while still only including full years, gives
mean(Ιn[1:mys]) - mean(Is[1:mys])

# which is a clear violation of Kepler's second law.

###################################################################### #src
###################################################################### #src
# # Timeseries analysis of reflected solar rad.
###################################################################### #src
###################################################################### #src

# %% Seasonal decomposition of outgoing rad #src
# ## Seasonal decomposition
fig, axs = subplots(2, 1; figsize = (figx, 8))

sinusoidal = Sinusoidal(Float32.([1, 2] ./ DAYS_IN_ORBIT))
Rn_sea, Rn_res = decompose(Rn, sinusoidal)
Rs_sea, Rs_res = decompose(Rs, sinusoidal)

seasonals = (Rn_sea, Rs_sea)
residuals = (Rn_res, Rs_res)

n, s = seasonals
g = (n .+ s) ./ 2
axs[1].plot(truetime[1:24], n[1:24], label = "\$\\mathcal{N}\$", marker = "o", ms = 5)
axs[1].plot(truetime[1:24], s[1:24]; ls = "--", label = "\$\\mathcal{S}\$", marker = "o", ms = 5)
axs[1].set_xlabel("days")

n, s = residuals
g = (n .+ s) ./ 2
axs[2].plot(t, n, label = "\$\\mathcal{N}\$")
axs[2].plot(t, s; ls = "--", label = "\$\\mathcal{S}\$")
axs[2].legend(loc="upper right", ncol = 3)
sn, = plottrend!(axs[2], t, n; color = "C0", addlabel=false)
ss, = plottrend!(axs[2], t, s; color = "C1", addlabel=false, ls = "dashed")

axs[1].set_ylabel("\$Y\$ (W m⁻²)")
axs[2].set_ylabel("\$R'\$ (W m⁻²)")
fig.tight_layout()
fig.align_ylabels(axs)
fig.subplots_adjust(hspace = 0.4, left = 0.13, bottom = 0.1)
wsave(papersdir("figures", "seasonal"), fig) #src

# Slopes of residuals
sn, ss

# Percentage of variance of R that is due to seasonal components:
var(Rn_sea)/var(Rn), var(Rs_sea)/var(Rs)

# Percentage of variance of R that is due to residual components:
var(Rn_res)/var(Rn), var(Rs_res)/var(Rs)
# Percentage of variance of R that is due to covariability of seasonal-residual
cov(Rn_sea, Rn_res)/var(Rn), cov(Rs_sea, Rs_res)/var(Rs)

# ## Residuals are noise
# %% #src
using Random, TimeseriesSurrogates, Distributions
fig = figure(figsize = (figx, 2figy/3))

ax_gaus = fig.add_subplot(1,2,1)
ax_mi = fig.add_subplot(1,2,2)

x = Rn_res # change to Rs_res for SH
# detrend before plotting distribution
_, _, trend = linreg(1:length(x), x)
xdet = x .- trend
# fit normal distribution
gaussian = Distributions.fit(Normal, xdet)
# q-q plot
q1 = quantile.(Ref(xdet), 0:0.01:1)
q2 = quantile.(gaussian, 0:0.01:1)
ax_gaus.plot(q1, q2; marker = "o", ls = "none")
ax_gaus.plot(-3:3, -3:3; lw = 1, color = "C1")
ax_gaus.set_xlim(-3, 3)
ax_gaus.set_ylim(-3, 3)
ax_gaus.set_xlabel("quantile res.")
ax_gaus.set_ylabel("quantile G.")

# ### Calculate mutual information
using Entropies
sg = surrogenerator(x, TFTS(0.1))

x = Array(Rn_res)[1:end-1]
y = Array(Rn_res)[2:end]
function mutualinfo(x, y, ε)
    Hx = genentropy(Dataset(x), ε)
    Hy = genentropy(Dataset(y), ε)
    Hxy = genentropy(Dataset(x,y), ε)
    return Hx + Hy - Hxy
end
bins = std(x)/4 # bins used for MI
mi = mutualinfo(x, y, bins)
M = 10000;
mis = zeros(M)

for i in 1:M
    local s = sg()
    local x = s[1:end-1]
    local y = s[2:end]
    mis[i] = mutualinfo(x, y, bins)
end

μ, σ = meanstd(mis)
ax_mi.axvspan(μ-2σ, μ+2σ, color = "C2", alpha = 0.25)
ax_mi.hist(mis, 50, color = "C3", density = true, label = "surrogate")
ax_mi.set_yticks([])
ax_mi.grid(false)
ax_mi.axvline(mi, color = "C0", label = "real")
ax_mi.legend(loc = "lower right")
ax_mi.set_xlabel("self-MI")
ax_mi.set_ylabel("pdf")

fig.tight_layout()
add_identifiers!(fig)
fig.subplots_adjust(wspace = 0.2, hspace = 0.3, left = 0.14, bottom = 0.2, top =0.92)
wsave(papersdir("figures", "residuals_noise"), fig) #src
fig

# ## Distributions for ΔR values
# %% #src
using TimeseriesSurrogates, Random
fig = figure(;figsize = (figx, 3figy/2))
ax_pps = fig.add_subplot(3,1,1)
ax_ar = fig.add_subplot(3,1,2; sharex = ax_pps)
ax_difpps = fig.add_subplot(3,2,5)
ax_difar = fig.add_subplot(3,2,6)

# ### Pseudoperiodic surrogates
τ = 3 # delay time = when AC becomes zero
ρ = 0.5 # evaluated using `noiseradius` function
d = 3 # embedding dimension
pp = PseudoPeriodic(d, τ, ρ, false)
sgnh = surrogenerator(Array(Rn), pp)
sgsh = surrogenerator(Array(Rs), pp)

Random.seed!(5151525)
ax_pps.plot(t, Rn, label = "\$\\mathcal{N}(R)\$", marker = "o", ms = 4)
ax_pps.plot(t, sgnh(), color = "C3", ls = "dashed", label = "surrogate", marker = "o", ms = 4)
ax_pps.legend(ncol=2)
ax_pps.set_ylabel("W m⁻²")

M = 5000
differences = zeros(M)
for i in 1:M
    rnh = sgnh()
    rsh = sgsh()
    differences[i] = timeagg(mean, t, rnh .- rsh)
end

μ, σ = meanstd(differences)
ax_difpps.axvspan(μ-2σ, μ+2σ, color = "C2", alpha = 0.25)
ax_difpps.hist(differences, 50, color = "C3", density = true, label = "surrogate")
ax_difpps.axvline(timemean(Rn .- Rs), color = "C0", label = "real")
ax_difpps.set_yticks([])
ax_difpps.grid(false)
ax_difpps.set_xticks(-0.25:0.25:0.25)
ax_difpps.set_ylabel("pdf")
ax_difpps.set_xlabel("avg. diff. \$\\Delta R\$ (W m⁻²)")
ax_difpps.legend(loc = "lower left")

# ### Autoregressive approximation
using StatsBase, ARFIMA, DSP, Random

D = Rn_res .- Rs_res
n = length(D)
σ = std(D)

# Now we estimate the parameters for the autoregressive process.
# There are several ways to do it. One can use the Yule-Walker equations, or
# hyper-optimization on a predefined length of parameters `p`. Here we will use
# linear predictive code estimation.

lpcmethod = DSP.LPCLevinson()
φ, e = DSP.lpc(D, 12, lpcmethod)
φ = -SVector(φ...) # definition of `φ` in ARFIMA is with -

rng = Random.MersenneTwister(612354152) # reproducibility
X = arma(rng, n+length(φ), σ, φ)[length(φ)+1:end]

ax_ar.plot(t, D, label = "\$D = R_\\mathcal{N}' - R_\\mathcal{S}'\$")
ax_ar.plot(t, X, label = "autoregressive", color = "C3")
ax_ar.set_xlim(tspan...)
ax_ar.legend(loc = "lower right", ncol = 2)
ax_ar.set_ylabel("W m⁻²")
ax_ar.set_yticks(-4:2:4)

N = 100000
X = arma(N, σ, φ) # simulate ultra long AR process
X .-= mean(X) # assert mean 0 for clarity
possible = Float64[]
for i in n:N-n
    push!(possible, mean(view(X, i:i+n)))
end
q1 = quantile(possible, 0.05)
q2 = quantile(possible, 0.95)
ax_difar.axvspan(q1, q2, color = "C2", alpha = 0.25)
ax_difar.hist(possible, color = "C3", 100, label = "autoregr.")
ax_difar.axvline(0.1, color = "C0", label = "real")
ax_difar.legend(loc="lower left")
ax_difar.set_yticks([])
ax_difar.set_xticks(-0.25:0.25:0.25)
ax_difar.grid(false)
ax_difar.set_xlabel("avg. diff. \$\\Delta R\$ (W m⁻²)")

fig.tight_layout()
fig.subplots_adjust(wspace = 0.1, hspace = 0.3, left = 0.12, bottom = 0.1, top =0.95)
add_identifiers!(fig)
wsave(papersdir("figures", "difference_is_zero"), fig) #src
fig

# ## Mean albedo decomposition + Fourier (only R)
# %% #src
fig, axs = subplots(2,1)

i = 1; X = R
Xn, Xs = hemispheric_means(zonalmean(X))
solarn, internaln = decompose(t, Xn.data, MeanAlbedo(Ιn.data))
solars, internals = decompose(t, Xs.data, MeanAlbedo(Is.data))

αn = timemean(Rn)/timemean(Ιn)
αs = timemean(Rs)/timemean(Is)

axs[i].plot(t, internaln, label = "\$\\mathcal{N}\$")
axs[i].plot(t, internals; ls = "--", label = "\$\\mathcal{S}\$")

in, is = var(internaln)/var(Xn), var(internals)/var(Xs)
println("Percentage of variance of internal for $(X.name)")
@show in, is

cn, cs = 2cov(solarn, internaln)/var(Xn), 2cov(solars, internals)/var(Xs)
println("Percentage of variance of covariability with sun for $(X.name)")
@show cn, cs

sn, ss = var(solarn)/var(Xn), var(solars)/var(Xs)
println("Percentage of variance of covariability with sun for $(X.name)")
@show sn, ss
println()

axs[1].legend(loc="upper right"; ncol=1)
axs[1].set_ylabel("\$R_\\mathrm{internal}\$")
axs[1].set_yticks(-5:5:10)
axs[1].set_xlim(Date(2008), Date(2015))

method = :FFTW
for (j, r) in enumerate((internaln[1:mys], internals[1:mys]))
    ## This can be made into Lomb-Scargle periodogram for more accuracy
    ## since the time axis is non-equi-spaced.
    if method == :FFTW
        P = abs2.(rfft(r .- mean(r)))
        f = FFTW.rfftfreq(length(r))
        ps = 1 ./ Float64.(f)
    elseif method == :LombScargle
        truetime = daysinmonth.(t[1:mys])
        truetime = Float32.(cumsum(truetime))
        periodogram = lombscargle(truetime, Array(r))
        ps, P = periodpower(periodogram)
        ps ./= 30 # divide by number of days in month (approx.)
    end
    spi = sortperm(ps)
    ps = ps[spi]; P = P[spi]
    ## Limit periods shown to `maxmonth`
    maxmonth = 2*12 + 2
    ℓ, k = findnearest.((2, maxmonth), Ref(ps))
    ps = ps[ℓ:k]
    P = P[ℓ:k]
    P = abs.(P)
    axs[2].plot(ps, log10.(P),
    color = "C$(j-1)", ls = j == 1 ? "-" : "--")
end

axs[2].set_ylabel("\$\\log_{10}(|\\mathcal{F}|)\$")
axs[2].set_xlabel("period (months)")

fig.tight_layout()
fig.subplots_adjust(left = 0.12, bottom = 0.12, hspace=0.22, top = 0.96)
wsave(papersdir("figures", "internal_new"), fig) #src
fig

# ## Longitude split (quadrant of hemisphere)
# %% #src
λs = dims(R, Lon)[1:5:size(R, Lon)]
d = zeros(length(λs), 4)
X, Y = hemispheric_functions(R)
for (i, λ) in enumerate(λs)
    X1, X2 = split_at_longitude(X, λ)
    Y1, Y2 = split_at_longitude(Y, λ)
    NE, NW = timemean.(spacemean.((X1, X2)))
    SE, SW = timemean.(spacemean.((Y1, Y2)))

    d[i, 1] = NE - NW
    d[i, 2] = NE - SE
    d[i, 3] = NE - SW
    d[i, 4] = SE - SW
end

fig = figure()
ax = gca()
for (i, n) in zip(1:4, ("NE - NW", "NE - SE", "NE - SW", "SE - SW", "NW - SW", "SE - SW"))
    z = d[:, i]
    a = mean(abs.(z))
    ax.plot(λs, z, label = "$(n), a = $(rdspl(a))")

end

ax.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc="lower left",
          ncol=2, mode="expand", borderaxespad=0.)
ax.set_xlabel("split longitude \$\\lambda\$ (ᵒ)")
ax.set_ylabel("avg. diff. (W m⁻²)")
ax.set_yticks(-6:3:6)
fig.tight_layout()

wsave(papersdir("figures", "longitude_split"), fig) #src

###################################################################### #src
###################################################################### #src
# # Cloudiness
###################################################################### #src
###################################################################### #src

# %% Compare the two different measures of Cloudiness on a MAP #src
# ## Definition and comparison

Tcorrected = sinusoidal_continuation(T, [1, 2.0]; Tmin = 0)
C1 = cloud_effective_albedo(F, Tcorrected, 0.9)
C1_timeavg = timemean(C1)

# use `\:arrow_down:` and tab for arrows
F_s_⬆ = ncread(EBAF_SFC, "sfc_sw_up_all_mon")
F_s_⬇ = ncread(EBAF_SFC, "sfc_sw_down_all_mon")
F_s_⬆_K = ncread(EBAF_SFC, "sfc_sw_up_clr_c_mon")
F_s_⬇_K = ncread(EBAF_SFC, "sfc_sw_down_clr_c_mon");

# Because surface data do not have the same amount of points, we are forced to use a
# smaller time span. (even if we time-average fields, this is still the accurate thing to do)
l = size(F_s_⬆, Time)
argsall = timemean.((I[Time(1:l)], R[Time(1:l)], F_s_⬆, F_s_⬇))
argsclr = timemean.((I[Time(1:l)], K[Time(1:l)], F_s_⬆_K, F_s_⬇_K))
α_ATM, α_SFC = surface_atmosphere_contributions(argsall...) # all sky
α_ATM_K, α_K_SFC = surface_atmosphere_contributions(argsclr...) # clear sky

C2_timeavg = α_ATM .- α_ATM_K

if false # plot spatial fields
    cmap = "gist_earth"
    as = (α_ATM, α_ATM_K, C2_timeavg, C1_timeavg, C1_timeavg .- C2_timeavg, α_ATM .+ α_SFC)
    nmes = ("α_ATM", "α_ATM_K", "C₂", "C₁", "C₁ - C₂", "α_TOA")
    kwargs = (vmin = 0, levels = 21, vmax = 0.7, cmap = "gist_earth")
    for (a, name) in zip(as, nmes)
        @show name, extrema(a)
        if name != "C₁ - C₂"
            fig, ax, = earthsurface(a; cmap = cmap, vmin = 0, vmax = 0.7)
        else
            fig, ax, = earthsurface(a; cmap = cmap, vmin = -0.2, vmax = 0.2)
        end
        ax.set_title(name)
    end
end

# We now need to estimate `g` so that the two versions match in value.
# Ιn the paper we call this the energetically consistent modification.
a = C2_timeavg; τ = timemean(Tcorrected); f = timemean(F)
_x = -2a ./ (τ .* a - f .* τ)
g = clamp.(1 .- _x ./ √3, 0, 1)

C = cloud_effective_albedo(F, Tcorrected, g)

# Let's produce maps of f, τ, g, C (for appendix)
for (label, X) in zip(("f", "\\tau", "g", "C"), (f, τ, g, timemean(C)))
    m = spacemean(X)
    fig, ax, = earthsurface(X; cmap = "gist_earth")
    ax.set_title("\$$(label)\$, mean = $(rdspl(m))")
    wsave(papersdir("figures", "map_$label"), fig) #src
end

# Average cloud albedo per hemisphere:
Cn, Cs = hemispheric_means(C)
timemean.((Cn, Cs))

# ## Plot of zonally average cloudiness, along with average values
# %% #src
fig = figure(figsize = (figx, 5))
Cz = timemean(zonalmean(C))
ax = gca()
ax.plot(sind.(dims(Cz, Lat)), Cz, label = "\$\\mathcal{T}(C)\$")
Fz = timemean(zonalmean(F))
Fz = (Fz .- mean(Fz)) ./ std(Fz)
Fz = Fz .* std(Cz) .+ mean(Cz)
ax.plot(sind.(dims(Cz, Lat)), Fz, ls = "dashed", color = "C2", label = "\$\\mathcal{T}(f)\$")
ax.legend(ncol = 2)

lats = dims(Cz, Lat) |> Array

latr = [-90, -30, 0, 30, 90]
s = 0.0
for i in 1:4
    l1, l2 = latr[i], latr[i+1]
    mv = latmean(Cz[Lat(Between(l1, l2))])
    x = (sind(l1) + sind(l2))/2
    ax.text(x, 0.025, rdspl(mv, 3), ha = :center)
    global s += mv
    ax.axvspan(sind(l1), sind(l2); alpha = 0.1, color = "C$(i+1)")
end
ax.set_ylim(0, 0.4)
llats = -70:20:70
ax.set_xticks(sind.(llats))
ax.set_xticklabels(llats)
ax.set_xlim(-1,1)
ax.set_xlabel("latitude");

fig.tight_layout()
fig.subplots_adjust(left = 0.1, bottom = 0.2, top = 0.95)
wsave(papersdir("figures", "cloud_lat"), fig) #src

# ## Plot of cloudiness timeseries and correlation with solar cycle
# %% #src

fig, axs = subplots(2,1; figsize = (figx, 7))
ax = axs[1]
ax.plot(t, Cn, label = "\$\\mathcal{N}(C)\$")
ax.plot(t, Cs; ls = "dashed", label = "\$\\mathcal{S}(C)\$")
ax.legend(loc = "upper right")
ax.set_yticks(0.14:0.03:0.20)
ax.set_xlim(tspan...)
ax.set_ylabel("\$C\$")

Cx = (Cn, Cs)
Ix = (Ιn, Is)
ax = axs[2]
lags = -13:13

monthlpv = Sinusoidal([1/12])
for j in 1:2
    cc = crosscor(Cx[j], Ix[j], lags)
    decompose(float.(lags), cc, monthlpv)
    φ = mod2pi(monthlpv.φ[2])
    φ = 12φ/2π # in months
    ax.plot(lags, cc, ls = j==1 ? "-" : "--",
    label="\$\\mathcal{$(j==1 ? "N" : "S")}, d \\approx $(rdspl(φ, 2))\$")
end
ax.set_xlabel("delay (months)"); ax.set_xticks(-12:3:12);
ax.legend(loc = "lower right")
ax.set_ylabel("crosscor.")

fig.tight_layout()
fig.subplots_adjust(left = 0.13, bottom = 0.15, top = 0.95)
add_identifiers!(fig)
fig.align_ylabels()
wsave(papersdir("figures", "cloud_and_sun"), fig) #src

# Weighted integrals of sunlight and cloudiness
D_with_sun = timeagg(mean, Cs, Is) - timeagg(mean, Cn, Ιn)
G_with_sun = 0.5(timeagg(mean, Cs, Is) + timeagg(mean, Cn, Ιn))
P_with_sun = 100 * D_with_sun/G_with_sun # proportion in percentage

D_no_sun = timeagg(mean, Cs) - timeagg(mean, Cn)
G_no_sun = 0.5(timeagg(mean, Cs) + timeagg(mean, Cn))
P_no_sun = 100 * D_no_sun/G_no_sun # proportion in percentage

P_with_sun, P_no_sun

# Alright, we now quantify the extra reflection one gets from the time delay
# by shifting the NH timeseries until to matches the sun.
using Interpolations
cloud_reflection_n = timemean(Cn .* Ιn)
cloud_reflection_s = timemean(Cs .* Is)
t = dims(Ιn, Time) |> Array
interp = CubicSplineInterpolation(1:length(t), Cn.data; extrapolation_bc = Line())

result = []
ds = 0:0.1:1.5
for d in ds
    r = timemean([interp(i) for i in (1:length(t)) .+ d] .* Ιn)
    push!(result, r)
end

fig = figure()
plot(ds, result, label = "Cloud reflection of NH")
axhline(cloud_reflection_n, ls = "dashed", label = "measured NH reflection", color = "C3")
axvline(1.2, color = "C1", label = "SH delay")
axvline(1.5, color = "C2", label = "0 delay")
ylabel("cloud reflection")
xlabel("delay shift (months)")
legend()
tight_layout()

# The effect of shifting the timeseries of NH by 1.2 months (so that they have same
# delay as SH) compensates the asymmetry by:
effect = result[13] - result[1]

# ## Cloudiness over specific land type
# %% #src
# Because `O, L, E` have a shorter timespan than `C`, we need to modify `C` to be
# in the same time span.

function hemispheric_weighted_avgs(X, fields)
    @assert size(X) == size(fields[1])
    Xn, Xs = hemispheric_functions(X)
    Xnt, Xst = timemean.((Xn, Xs))
    northpie = []
    southpie = []
    northcov = []
    southcov = []

    for W in fields
        Wn, Ws = hemispheric_functions(W)
        Gn, Gs = spacemean(Xn, Wn), spacemean(Xs, Ws)
        Gna, Gsa = timemean.((Gn, Gs))
        Wnt, Wst = timemean.((Wn, Ws))
        push!(northpie, (Gna, spacemean(Wnt)))
        push!(southpie, (Gsa, spacemean(Wst)))
        ## calcualte no covariability
        push!(northcov, spacemean(Xnt, Wnt))
        push!(southcov, spacemean(Xst, Wst))
    end
    return northpie, southpie, northcov, southcov
end

ol = size(O, Time)
northpie, southpie, northcov, southcov = hemispheric_weighted_avgs(C[Time(1:ol)], (O, E, L));

# Put the numbers into bar plot
# %% Barplot over surface #src
fig = figure()
ax = fig.add_subplot()
colors = ["C2", "C3", "C0"]
fig.tight_layout()
prefix = 1.2

for (i, data) in enumerate((northpie, southpie))
    pre = i == 1 ? 0.0 : prefix
    widths = [s[2] for s in data]
    x = [0, cumsum(widths)[1:end-1]...] .+ pre
    h = [s[1] for s in data]
    hcov = (northcov, southcov)[i]

    ax.bar(x, h; width = widths,
    color = colors, align = "edge")
    ax.bar(x, hcov; width = widths, linewidth = 3.0,
    color = colors, align = "edge", ls = "dashed", fill = false, edgecolor = lighten.(colors, 0.55))
    pos = x .+ widths./2
    for (j, (c, w)) in enumerate(data)
        str = "$(rdspl(c, 2))\n$(rdspl(w, 2))"
        off = + 0.015
        i == 2 && j == 2 && (off += 0.01)
        ax.text(pos[j], h[j] + off, string(rdspl(c, 2)), ha = :center, color = colors[j])
        botoff = (i == 2 && j == 2) ? +0.01 : -0.02
        ax.text(pos[j], botoff, string(rdspl(w, 2)), ha = :center)
    end
end
ax.spines["right"].set_visible(false)
ax.spines["top"].set_visible(false)
ax.set_xticks([])
ax.spines["bottom"].set_visible(false)

ax.set_ylim(0, 0.25)
ax.set_ylabel("\$\\langle C \\rangle\$ over surface")
ax.text(0.5, 0.25, "northern", ha = :center)
ax.text(0.5 + prefix, 0.25, "southern", ha = :center)
fig.tight_layout()

## Create legend
mpatches = matplotlib.patches
handles = [mpatches.Patch(color=c, label=F) for (c, F) in zip(colors, (:ocean, :ice, :land))]

ax.legend(handles=handles; ncol=3, bbox_to_anchor=(0.0, -0.3, 1, .102), loc="lower left",
            mode="expand", borderaxespad=0.)

fig.tight_layout()
fig.subplots_adjust(bottom = 0.24)

wsave(papersdir("figures", "cloud_over_surface_bar"), fig) #src

# %% total albedo given surface and cloud #src
# ## Total albedo given surface and cloud albedo and transmittance
# Main goal here is to see impact of having cloud albedo difference of `0.02`.
transmitance = 0.53 # some given trasmittance

asfc = 0:0.01:1.0
aatm = 0:0.01:transmitance

achoice = [0.18, 0.2, 0.22]
schoice = [0.1, 0.3, 0.6]

fig = figure()
ax1 = fig.add_subplot(1, 2, 1)
atoa = total_toa_albedo.(aatm, asfc', transmitance)
ax1.pcolormesh(aatm, asfc, atoa')
ax1.set_title("transmittance = $(transmitance)")
ax2 = fig.add_subplot(2, 2, 2)
ax3 = fig.add_subplot(2, 2, 4)

for i in 1:length(achoice)
    a = achoice[i]
    ax1.axvline(a, color = "C$i", ls = "dashed")
    atotal = total_toa_albedo.(a, asfc, transmitance)
    ax2.plot(asfc, atotal, color = "C$i", label = "$(a)")
    ax2.set_xlabel("a_sfc")
    ax2.legend()
end

for i in 1:length(schoice)
    a = schoice[i]
    ax1.axhline(a, color = "C$i", ls = "dashed")
    atotal = total_toa_albedo.(aatm, a, transmitance)
    ax3.plot(aatm, atotal, color = "C$i", label = "$a")
    ax3.set_xlabel("a_atm")
    ax3.legend()
end

fig.tight_layout()

# So for a broad range of values (and you have to zoom into the figure), we find a
# total change of 0.02. This is also the surface abledo assymetry
Kn, Ks = hemispheric_means(K)
δαK = timemean(Kn ./ Ιn) - timemean(Ks ./ Is)

###################################################################### #src
###################################################################### #src
# # Appendices
###################################################################### #src
###################################################################### #src
# ## Continuation of value of τ
# %% #src
count(ismissing, T)/length(T) # about 6% of the points are missing

fig, axs = subplots(2, 1; figsize = (figx, 6), sharex = true)
for (i, l) in enumerate((-82, 82))
    ax = axs[i]
    example = T[Lat(Near(l)), Lon(Near(5))]
    example = Vector{Any}(Array(example))

    lpv = Sinusoidal(Float32.([1, 2] ./ DAYS_IN_ORBIT))
    nmi = findall(!ismissing, example)
    mi = findall(ismissing, example)
    x0 = Array{Float32}(example[nmi])
    t0 = truetime[nmi]
    sea, res = decompose(t0, x0, lpv)
    c = ClimateBase.cosines(truetime, lpv)
    acosine = copy(example)
    acosine[mi] .= c[mi]
    example[mi] .= nothing

    ax.plot(t, example, marker = "o", label = "lat = $(l)", ms = 4)
    ax.plot(t, c, ls = "dotted", alpha = 0.75, label = "cosine cont.")
    ax.legend(loc = "upper right")
    ax.set_ylabel("\$\\tau\$")
end
axs[2].set_xlim(tspan...)
fig.tight_layout()
fig.subplots_adjust(left = 0.1, right = 0.95, bottom = 0.14, hspace=0.1, top = 0.95)
wsave(papersdir("figures", "tau_continuation"), fig) #src

# Finally, make a Jupyter notebook out of the entire file. #src
using Literate #src
Literate.notebook(papersdir("figures.jl"), projectdir(); execute = false, credit = false) #src


### A Pluto.jl notebook ###
# v0.17.0

using Markdown
using InteractiveUtils

# ╔═╡ b787a8ca-7bd4-40a7-8574-334c72eda77d
begin
	using PrettyTables
	using PlutoUI
end

# ╔═╡ 67909d8a-421f-41cb-bc87-c1488be36074
md"""
### Practice Problem: Gas Phase Chemical Reaction Equilibrium Calculation

The gas-phase reaction:

$$A_{(g)}+B_{(g)} ~{\rightleftarrows}~ 2C_{(g)}$$

occurs in a closed well-mixed reactor with constant volume V = 10~L, a constant T = 500$^{\circ}$C, but with an unknown pressure (measured in bar). Initailly, there is 2 mol of A and 1 mol of B and no product C. Let the index of A = 1, B = 2 and C = 3. The Gibbs energies of formation at standard temperature and pressure are given by: $G_{1}^{\circ}$ = -2.0 kJ/mol,  $G_{2}^{\circ}$ = -3.6 kJ/mol and $G_{3}^{\circ}$ = -6.7 kJ/mol, respectively.

The teaching team measured the equilibrium extent of reaction $\epsilon^{\star}_{1}$ = 0.67 mol for these conditions. However, the _mythical mole meter_ used by the teaching team to measure this extent has been acting up (and may be giving the wrong reading).

__Assume__:
* The gas phase is _ideal_ for all components
* The ideal gas law is valid for this system
* The reference pressure $P^{\circ}$ = 1 bar.
* R = 8.314 J mol$^{-1}$ K$^{-1}$ and R = 0.08314 bar L mol$^{-1}$ K$^{-1}$.

__Question__:
* Is the measured equilibrium extent of reaction consistent with the problem setup?

"""

# ╔═╡ 80b572d7-3400-4373-af66-a394fe45296f
begin

	# setup the problem -
	V = 10.0 			# units: L
	T = (500+273.15)	# units: K
	R = 0.008314 		# units: kJ mol^-1 K^-1
	R_bar = 0.08314 	# units: bar L mol^-1 K^-1
	Pₒ = 1.0 			# units: bar

	# G_formation_array -
	G_formation_array = [
		-2.0 ; # 1 G_A units: kJ/mol
		-3.6 ; # 2 G_B units: kJ/mol
		-6.7 ; # 3 G_C units: kJ/mol
	];

	# reaction stoichiometry -
	S = [
		-1.0 ; 	# 1 A
		-1.0 ; 	# 2 B
		2.0 ; 	# 3 C
	];

	# initial mol -
	n_initial = [
		2.0 ; # 1 A units: mol
		1.0 ; # 2 B units: mol
		0.0 ; # 3 C units: mol
	];

	# what is the measured extent?
	ϵ₁ = 0.67 # units: mol

	# show -
	nothing
end

# ╔═╡ 9c7e21c6-dfe3-4f3e-b2e0-23d7ab578619
md"""
#### Solution: Compute the Eqilibrium constant $K$ using two different approaches and compare

To determine if the _mythical mole meter_ is correct, we can compute the equilibrium constant using the extent $\epsilon_{1}$, and then compare that value with the equilibrium constant computed using the $\Delta{G}^{\circ}$. If the values are the same, the _mythical mole meter_ is working, otherwise the _mythical mole meter_ is broken (and giving bad measurements of the extent).

"""

# ╔═╡ 4a63eb66-7090-41d0-80b3-94b56f5edebf
md"""
###### Compute K$_{\Delta{G}}$: equilibrium constant using the ΔG values
"""

# ╔═╡ 4e10a9b5-cb5e-4b4a-87b4-0ce024a1ad6f
begin

	# compute the ΔG -
	ΔG = sum(G_formation_array.*S)
	RT = R*T

	# compute the equlibrium constant using the ΔG -
	K_ΔG = exp(-ΔG/RT)

	# display -
	with_terminal() do
		println("K_ΔG = $(K_ΔG)")
	end
end

# ╔═╡ 9f253794-d773-48bc-a1b5-94ca242ff798
md"""
###### Compute K$_{\epsilon}$: equilibrium constant using the extent values

Let's compute the equilibrium constant from the extent of reaction. Setup the state table using the species mole balance:

$$n_{i} = n^{\circ}_{i} + \sigma_{i1}\epsilon_{1}$$

"""

# ╔═╡ 47e684a6-773b-4009-b7cc-23dc2fbf4a4b
begin

	# compute the final number of mols, mol fraction -
	n_final = n_initial .+ S*ϵ₁
	n_total = sum(n_final)
	y_final = (1/n_total)*n_final

	# setup data array =
	state_table_array = Array{Any,2}(undef, 3,5)

	# A -
	state_table_array[1,1] = "A"
	state_table_array[1,2] = S[1]
	state_table_array[1,3] = n_initial[1]
	state_table_array[1,4] = n_final[1]
	state_table_array[1,5] = y_final[1]

	# B -
	state_table_array[2,1] = "B"
	state_table_array[2,2] = S[2]
	state_table_array[2,3] = n_initial[2]
	state_table_array[2,4] = n_final[2]
	state_table_array[2,5] = y_final[2]

	# C -
	state_table_array[3,1] = "C"
	state_table_array[3,2] = S[3]
	state_table_array[3,3] = n_initial[3]
	state_table_array[3,4] = n_final[3]
	state_table_array[3,5] = y_final[3]

	# setup a pretty table =
	# header row -
	state_table_header_row = (["Species","σ_i1","nᵢ_initial","nᵢ_final","yᵢ_final"],["","","mol", "mol", ""]);

	with_terminal() do
		# write the table -
		pretty_table(state_table_array; header=state_table_header_row)
	end

end

# ╔═╡ 8dca4509-6d6b-41c2-9d22-81c4563a5136
md"""
Now that we have the state table, we can use the gas phase equlibrium expression:

$$\prod_{i=1}^{\mathcal{M}}\left(y_{i}\hat{\phi}_{i}^{v}\right)^{\sigma_{i1}} = \left(\frac{P}{P^{\circ}}\right)^{-\sigma_{1}}K$$

To compute the equlibrium constant using the extent. Becuase we have assumed ideality, all the fugacity coefficient terms $\hat{\phi}^{v}_{i}\rightarrow\forall{i}$. In addition the pressure term is 1 becuase $\sigma_{1} = 0$ (the number of moles is conserved).
Thus, the equilibrium constant is given by:

$$K = \frac{y_{3}^2}{y_{1}y_{2}}$$
"""

# ╔═╡ b7c489c1-513b-4c1c-b22a-9c0738e4a62e
begin

	# compute the overall stoichiometric coefficient -
	σ₁ = sum(S);

	# Compute pressure term -
	P = (1/V)*(R_bar*T)*n_total
	P_term = (P/Pₒ)^(-σ₁)

	# compute the K -
	K_ϵ = (y_final[1]^S[1])*(y_final[2]^S[2])*(y_final[3]^S[3])*(1/P_term)

	with_terminal() do
		println("K_ϵ = $(K_ϵ)")
	end
end

# ╔═╡ 9ce4fb80-a6da-4a3d-88f2-83092be1f340
md"""
### Bummer. K$_{\Delta{G}}$ $\neq$ K$_{\epsilon}$ $\longrightarrow$ the _mythical mole meter_ is broken!
"""

# ╔═╡ ccc364fe-3d4a-11ec-3415-0753629989e9
html"""
<style>
main {
    max-width: 1200px;
    width: 85%;
    margin: auto;
    font-family: "Roboto, monospace";
}

a {
    color: blue;
    text-decoration: none;
}

.H1 {
    padding: 0px 30px;
}
</style>"""

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"

[compat]
PlutoUI = "~0.7.18"
PrettyTables = "~1.2.3"
"""

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

[[AbstractPlutoDingetjes]]
deps = ["Pkg"]
git-tree-sha1 = "0ec322186e078db08ea3e7da5b8b2885c099b393"
uuid = "6e696c72-6542-2067-7265-42206c756150"
version = "1.1.0"

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

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

[[Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[Crayons]]
git-tree-sha1 = "3f71217b538d7aaee0b69ab47d9b7724ca8afa0d"
uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
version = "4.0.4"

[[DataAPI]]
git-tree-sha1 = "cc70b17275652eb47bc9e5f81635981f13cea5c8"
uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
version = "1.9.0"

[[DataValueInterfaces]]
git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6"
uuid = "e2d170a0-9d28-54be-80f0-106bbe20a464"
version = "1.0.0"

[[Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

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

[[Formatting]]
deps = ["Printf"]
git-tree-sha1 = "8339d61043228fdd3eb658d86c926cb282ae72a8"
uuid = "59287772-0a20-5a39-b81b-1366585eb4c0"
version = "0.4.2"

[[Hyperscript]]
deps = ["Test"]
git-tree-sha1 = "8d511d5b81240fc8e6802386302675bdf47737b9"
uuid = "47d2ed2b-36de-50cf-bf87-49c2cf4b8b91"
version = "0.0.4"

[[HypertextLiteral]]
git-tree-sha1 = "5efcf53d798efede8fee5b2c8b09284be359bf24"
uuid = "ac1192a8-f4b3-4bfe-ba22-af5b92cd3ab2"
version = "0.9.2"

[[IOCapture]]
deps = ["Logging", "Random"]
git-tree-sha1 = "f7be53659ab06ddc986428d3a9dcc95f6fa6705a"
uuid = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
version = "0.2.2"

[[InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[IteratorInterfaceExtensions]]
git-tree-sha1 = "a3f24677c21f5bbe9d2a714f95dcd58337fb2856"
uuid = "82899510-4779-5014-852e-03e436cf321d"
version = "1.0.0"

[[JSON]]
deps = ["Dates", "Mmap", "Parsers", "Unicode"]
git-tree-sha1 = "8076680b162ada2a031f707ac7b4953e30667a37"
uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
version = "0.21.2"

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

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

[[LibGit2]]
deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

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

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

[[LinearAlgebra]]
deps = ["Libdl"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

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

[[Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

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

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

[[Parsers]]
deps = ["Dates"]
git-tree-sha1 = "ae4bbcadb2906ccc085cf52ac286dc1377dceccc"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.1.2"

[[Pkg]]
deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

[[PlutoUI]]
deps = ["AbstractPlutoDingetjes", "Base64", "Dates", "Hyperscript", "HypertextLiteral", "IOCapture", "InteractiveUtils", "JSON", "Logging", "Markdown", "Random", "Reexport", "UUIDs"]
git-tree-sha1 = "57312c7ecad39566319ccf5aa717a20788eb8c1f"
uuid = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
version = "0.7.18"

[[PrettyTables]]
deps = ["Crayons", "Formatting", "Markdown", "Reexport", "Tables"]
git-tree-sha1 = "d940010be611ee9d67064fe559edbb305f8cc0eb"
uuid = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
version = "1.2.3"

[[Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

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

[[Reexport]]
git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "1.2.2"

[[SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

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

[[Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[TOML]]
deps = ["Dates"]
uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"

[[TableTraits]]
deps = ["IteratorInterfaceExtensions"]
git-tree-sha1 = "c06b2f539df1c6efa794486abfb6ed2022561a39"
uuid = "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c"
version = "1.0.1"

[[Tables]]
deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "LinearAlgebra", "TableTraits", "Test"]
git-tree-sha1 = "fed34d0e71b91734bf0a7e10eb1bb05296ddbcd0"
uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
version = "1.6.0"

[[Tar]]
deps = ["ArgTools", "SHA"]
uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"

[[Test]]
deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

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

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

[[p7zip_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
"""

# ╔═╡ Cell order:
# ╟─67909d8a-421f-41cb-bc87-c1488be36074
# ╠═80b572d7-3400-4373-af66-a394fe45296f
# ╟─9c7e21c6-dfe3-4f3e-b2e0-23d7ab578619
# ╟─4a63eb66-7090-41d0-80b3-94b56f5edebf
# ╠═4e10a9b5-cb5e-4b4a-87b4-0ce024a1ad6f
# ╟─9f253794-d773-48bc-a1b5-94ca242ff798
# ╠═47e684a6-773b-4009-b7cc-23dc2fbf4a4b
# ╠═8dca4509-6d6b-41c2-9d22-81c4563a5136
# ╠═b7c489c1-513b-4c1c-b22a-9c0738e4a62e
# ╟─9ce4fb80-a6da-4a3d-88f2-83092be1f340
# ╠═b787a8ca-7bd4-40a7-8574-334c72eda77d
# ╟─ccc364fe-3d4a-11ec-3415-0753629989e9
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


struct FenwickTree{T}
    bi_tree::Vector{T} #bi_tree is shorthand for Binary Indexed Tree, an alternative name for Fenwick Tree
    n::Int
end

"""
    FenwickTree{T}(n::Integer)

Constructs a [`FenwickTree`](https://en.wikipedia.org/wiki/Fenwick_tree) of length `n`.

"""
FenwickTree{T}(n::Integer) where T = FenwickTree{T}(zeros(T, n), n)

"""
    FenwickTree(counts::AbstractArray)

Constructs a [`FenwickTree`](https://en.wikipedia.org/wiki/Fenwick_tree) from an array of `counts`

"""
function FenwickTree(a::AbstractVector{U}) where U
    n = length(a)
    tree = FenwickTree{U}(n)
    @inbounds for i = 1:n
        inc!(tree, i, a[i])
    end
    tree
end

length(ft::FenwickTree) = ft.n
Base.eltype(::Type{FenwickTree{T}}) where T = T

"""
    inc!(ft::FenwickTree{T}, ind::Integer, val)

Increases the value of the [`FenwickTree`] by `val` from the index `ind` upto the length of the Fenwick Tree.

"""
function inc!(ft::FenwickTree{T}, ind::Integer, val = 1) where T
    val0 = convert(T, val)
    i = ind
    n = ft.n
    @boundscheck 1 <= i <= n || throw(ArgumentError("$i should be in between 1 and $n"))
    @inbounds while i <= n
        ft.bi_tree[i] += val0
        i += i&(-i)
    end
end

"""
    dec!(ft::FenwickTree, ind::Integer, val)

Decreases the value of the [`FenwickTree`] by `val` from the index `ind` upto the length of the Fenwick Tree.

"""
dec!(ft::FenwickTree, ind::Integer, val = 1 ) = inc!(ft, ind, -val)

"""
    incdec!(ft::FenwickTree{T}, left::Integer, right::Integer, val)

Increases the value of the [`FenwickTree`] by `val` from the indices from `left` and decreases it from the `right`.

"""
function incdec!(ft::FenwickTree{T}, left::Integer, right::Integer, val = one(T)) where T
    val0 = convert(T, val)
    inc!(ft, left, val0)
    dec!(ft, right, val0)
end

"""
    prefixsum(ft::FenwickTree{T}, ind::Integer)

Return the cumulative sum from index 1 upto `ind` of the [`FenwickTree`](@ref)

# Examples
```
julia> f = FenwickTree{Int}(6)
julia> inc!(f, 2, 5)
julia> prefixsum(f, 1)
 0
julia> prefixsum(f, 3)
 5
```
"""
function prefixsum(ft::FenwickTree{T}, ind::Integer) where T
    sum = zero(T)
    i = ind
    n = ft.n
    @boundscheck 1 <= i <= n || throw(ArgumentError("$i should be in between 1 and $n"))
    @inbounds while i > 0
        sum += ft.bi_tree[i]
        i -= i&(-i)
    end
    sum
end

getindex(ft::FenwickTree{T}, ind::Integer) where T = prefixsum(ft, ind)


abstract type ActivationFunction end

struct Linear <: ActivationFunction end
(::Linear)(x::T) where {T <: Number} = x

struct ReLU <: ActivationFunction end
(::ReLU)(x::T) where {T <: Number} = max(zero(x), x)

struct Sigmoid <: ActivationFunction end
(::Sigmoid)(x::T) where {T <: Number} = 1 / (1 + exp(-x))

struct Square <: ActivationFunction end
(::Square)(x::T) where {T <: Number} = x^2

struct Tanh <: ActivationFunction end
(::Tanh)(x::T) where {T <: Number} = tanh(x)


"""
    ScaledKernel(k::Kernel, σ²::Real=1.0)

Scaled kernel derived from `k` by multiplication with variance `σ²`.

# Definition

For inputs ``x, x'``, the scaled kernel ``\\widetilde{k}`` derived from kernel ``k`` by
multiplication with variance ``\\sigma^2 > 0`` is defined as
```math
\\widetilde{k}(x, x'; k, \\sigma^2) = \\sigma^2 k(x, x').
```
"""
struct ScaledKernel{Tk<:Kernel,Tσ²<:Real} <: Kernel
    kernel::Tk
    σ²::Vector{Tσ²}
end

function ScaledKernel(kernel::Tk, σ²::Tσ²=1.0) where {Tk<:Kernel,Tσ²<:Real}
    @check_args(ScaledKernel, σ², σ² > zero(Tσ²), "σ² > 0")
    return ScaledKernel{Tk,Tσ²}(kernel, [σ²])
end

@functor ScaledKernel

(k::ScaledKernel)(x, y) = first(k.σ²) * k.kernel(x, y)

function kernelmatrix(κ::ScaledKernel, x::AbstractVector, y::AbstractVector)
    return κ.σ² .* kernelmatrix(κ.kernel, x, y)
end

function kernelmatrix(κ::ScaledKernel, x::AbstractVector)
    return κ.σ² .* kernelmatrix(κ.kernel, x)
end

function kernelmatrix_diag(κ::ScaledKernel, x::AbstractVector)
    return κ.σ² .* kernelmatrix_diag(κ.kernel, x)
end

function kernelmatrix_diag(κ::ScaledKernel, x::AbstractVector, y::AbstractVector)
    return κ.σ² .* kernelmatrix_diag(κ.kernel, x, y)
end

function kernelmatrix!(
    K::AbstractMatrix, κ::ScaledKernel, x::AbstractVector, y::AbstractVector
)
    kernelmatrix!(K, κ.kernel, x, y)
    K .*= κ.σ²
    return K
end

function kernelmatrix!(K::AbstractMatrix, κ::ScaledKernel, x::AbstractVector)
    kernelmatrix!(K, κ.kernel, x)
    K .*= κ.σ²
    return K
end

function kernelmatrix_diag!(K::AbstractVector, κ::ScaledKernel, x::AbstractVector)
    kernelmatrix_diag!(K, κ.kernel, x)
    K .*= κ.σ²
    return K
end

Base.:*(w::Real, k::Kernel) = ScaledKernel(k, w)

Base.show(io::IO, κ::ScaledKernel) = printshifted(io, κ, 0)

function printshifted(io::IO, κ::ScaledKernel, shift::Int)
    printshifted(io, κ.kernel, shift)
    return print(io, "\n" * ("\t"^(shift + 1)) * "- σ² = $(first(κ.σ²))")
end


#states constructors
struct State{T<:AbstractVector}
    st::T
    ope::Op
    State(st,ope) = log(2,length(st)) != dim(ope) ? error("lenght of vector does not match dimension") : new{typeof(st)}(st/sqrt(st'*st),ope)
end

st(s::State) = s.st

ope(s::State) = s.ope

dim(s::State) = dim(s.ope)

basis(s::State) = basis(s.ope)

typ(s::State) = eltype(s.st)

function rhosp(sta::State, prec=15)
    n = dim(ope(sta))
    rhospv = spzeros(typ(sta),n,n)
    estate = st(sta)
    o = ope(sta)
    for i in 1:n
        for j in 1:n
            rhospv[i,j] = round(estate'*ada(o, i, j)*estate, digits = prec)
        end
    end
    return rhospv
end

function kqsp(sta::State, prec=15)
    n = dim(ope(sta))
    k = spzeros(typ(sta),n,n)
    estate = st(sta)
    o = ope(sta)
    for i in 1:n
        for j in 1:n
            k[i,j] = round(estate'*aa(o, j, i)*estate, digits = prec)
        end
    end
    return k
end

rhoqsp(s::State, prec=15) = [rhosp(s, prec) kqsp(s, prec); kqsp(s, prec)' I-rhosp(s, prec)']

function non_zero(c, prec=15)
    a = similar(c, Int)
    cof = similar(c, Float64)
    count = 1
    @inbounds for i in eachindex(c)
        a[count] = i
        cof[count] = round(c[i], digits = prec)
        count += (round(c[i], digits = prec) != zero(eltype(c)))
    end
    return resize!(a, count-1), resize!(cof, count-1)
end

#states constructors for both states_fixed and states
#disadvantage is that working with unmutables is less efficient
#also I believe it would be better to have both structs
# mutable struct State{T<:AbstractVector}
#     st::T
#     ope::Op
#     nume::Union{Int64,Missing}
# end
#
# function State(st,ope;nume=missing)
#     return State(st,ope,nume)
# end

# non_zero inputs a vector c and a precision (number of digits)
# and outputs two vectors: the first one indicating the indexes
# of non zero elements (up to precision prec) and the second
# the coefficient in each of these non zero indexes


consensus(u, v) = (isequal(u.x,v.x) || v.x == nothing ? u.x : (u.x == nothing ? v.x : nothing))

function Base.:+(u::Poly, v::Poly)
    x = consensus(u, v)
    x === nothing && return poly(sym(u) + sym(v), nothing)
    q = Poly(u.T, x)

    for k = 0:max(u.n, v.n)
        q[k] = u[k] + v[k]
    end
    q
end

Base.:+(u, v::Poly) = polyas(v, u) + v
Base.:+(u::Poly, v) = u + polyas(u, v)

function Base.:-(u::Poly, v::Poly)
    x = consensus(u, v)
    x === nothing && return poly(sym(u) - sym(v), nothing)
    q = Poly(u.T, x)

    for k = 0:max(u.n, v.n)
        q[k] = u[k] - v[k]
    end
    q
end

Base.:-(u::Poly, v) = u - polyas(u, v)
Base.:-(u, v::Poly) = polyas(v, u) - v

function Base.:*(u::Poly, v::Poly)
    x = consensus(u, v)
    x === nothing && return poly(sym(u) * sym(v), nothing)
    q = Poly(u.T, x)

    for i = 0:u.n
        for j = 0:v.n
            q[i+j] += u[i] * v[j]
        end
    end
    q
end

Base.:*(u::Poly, v) = u * polyas(u, v)
Base.:*(u, v::Poly) = polyas(v, u) * v

function divide_poly(u::Poly, v::Poly; op=/)
    x = consensus(u, v)
    x === nothing && return poly(sym(u) / sym(v), nothing)
    q = Poly(u.T, x)
    r = copy(u)
    m, n = u.n, v.n

    for i = m-n:-1:0
        q[i] = op(r[i+n], v[n])
        for j = n+i:-1:i
            r[j] = r[j] - q[i] * v[j-i]
        end
    end

    update_order!(r)

    q, r
end

function Base.:/(u::Poly, v::Poly)
    q, r = divide_poly(u, v)
    q
end

Base.:/(u::Poly, v) = u / polyas(u, v)
Base.:/(u, v::Poly) = polyas(v, u) / v

function Base.:%(u::Poly, v::Poly)
    q, r = divide_poly(u, v)
    r
end

Base.:%(u::Poly, v) = u % polyas(u, v)
Base.:%(u, v::Poly) = polyas(v, u) % v

function pseudo_divide_poly(u::Poly, v::Poly)
    m, n = u.n, v.n
    α = v[n]^(m-n+1)
    q, r = divide_poly(u*α, v; op=÷)
    return q, r, α
end

Base.:^(p::Poly, k) = poly(sym(p)^k, p.x)

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

function gcd_naive(u::Poly, v::Poly)
    while degree(v) > 0
        println(u, '\t', v)
        u, v = v, u % v
    end
    return iszero(v) ? u : 1
end

function gcd_monic(u::Poly, v::Poly)
    if degree(u) == 0 || degree(v) == 0
        return gcd(cont(u), cont(v))
    end
    u = u / leading(u)
    while degree(v) > 0
        v = v / leading(v)
        u, v = v, u % v
    end
    return iszero(v) ? u : 1
end

function gcd_extended(u::Poly, v::Poly; op=/)
    if degree(u) == 0 || degree(v) == 0
        return gcdx(cont(u), cont(v))
    end

    sᵤ = one(u.T)
    tᵤ = zero(u.T)
    sᵥ = zero(v.T)
    tᵥ = one(v.T)

    while degree(v) > 0
        q, r = divide_poly(u, v; op)
        s, t = sᵤ - q*sᵥ, tᵤ - q*tᵥ
        u, sᵤ, tᵤ = v, sᵥ, tᵥ
        v, sᵥ, tᵥ = r, s, t
    end

    l = leading(u)
    if iszero(l)
        return u, sᵤ, tᵤ
    else
        return u / l, sᵤ / l, tᵤ / l
    end
end

function Base.gcd(u::Poly, v::Poly)
    if u.T <: Rational && u.T <: Rational
        return gcd_monic(u, v)
    else
        return gcd_naive(u, v)
    end
end

Base.gcd(u::Poly, v) = gcd(u, polyas(u, v))
Base.gcd(u, v::Poly) = gcd(polyas(v, u), v)

Base.gcdx(u::Poly, v::Poly) = gcd_extended(rationalize(u), rationalize(v))
Base.gcdx(u::Poly, v) = gcdx(u, polyas(u, v))
Base.gcdx(u, v::Poly) = gcdx(polyas(v, u), v)

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

Symbolics.derivative(p::Poly, x) = isequal(p.x, x) ? derivative(p) : derivative(sym(p), x)

function Symbolics.derivative(p::Poly)
    q = Poly(p.T, p.x)
    for i = 1:p.n
        q[i-1] = i * p[i]
    end
    q
end

function integrate(p::Poly)
    q = Poly(p.T, p.x)
    for i = 0:p.n
        q[i+1] = p[i] / (i+1)
    end
    q
end

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

function rationalize(p::Poly)
    p.T <: Rational && return p
    try
        q = Poly(Rational{BigInt}, p.x, p.n, Rational{BigInt}(p.a₀), Dict())
        for k = 1:p.n
            q[k] = Rational{BigInt}(p[k])
        end
        return q
    catch
        return nothing
    end
end

function to_monic(p::Poly)
    p = rationalize(p)
    c = leading(p)
    n = degree(p)
    q = Poly(p.T, p.x)
    for i = 0:n
        q[i] = p[i] * c ^ (n-1-i)
    end
    c, q
end

function from_monic(p::Poly, c)
    q = Poly(p.T, p.x)
    n = degree(p)
    for i = 0:n
        q[i] = p[i] / c ^ (n-1-i)
    end
    q
end

function integer_poly(p::Poly)
    p = rationalize(p)
    l = lcm(denominator.(extract_coef(p))...)
    return 1//l, p*l
end


module CoordinateDescent

using ProximalBase
using DataStructures: nlargest
using SparseArrays
using Statistics
using LinearAlgebra

export
  lasso, sqrtLasso, feasibleLasso!, scaledLasso, scaledLasso!,
  LassoPath, refitLassoPath,
  IterLassoOptions,

  # CD
  CoordinateDifferentiableFunction,
  CDLeastSquaresLoss, CDWeightedLSLoss, CDQuadraticLoss, CDSqrtLassoLoss,
  CDOptions,
  coordinateDescent, coordinateDescent!,

  # var coef
  GaussianKernel, SmoothingKernel, EpanechnikovKernel, evaluate, createKernel,
  locpoly, locpolyl1

include("utils.jl")
include("atom_iterator.jl")
include("cd_differentiable_function.jl")
include("coordinate_descent.jl")
include("lasso.jl")
include("varying_coefficient_lasso.jl")

end


using ForwardDiffOverMeasurements
using Measurements
using ForwardDiff
using Test

f1(x,y) = x+y
f2(x,y) = x-y
f3(x,y) = x*y
f4(x,y) = x/y

@testset "ForwardDiffOverMeasurements.jl" begin
    x1 = 1.0 ± 0.1
    y1 = 2.0 ± 0.001
    for op in (:f1,:f2,:f3,:f4)
        @eval begin
            @test ForwardDiff.derivative(x->$(op)(x,$y1),$x1) isa Measurement
            @test ForwardDiff.derivative(y->$(op)($x1,y),$y1) isa Measurement
        end
    end
end


module FastGaussQuadrature

using Compat

export gausslegendre
export gausschebyshev
export gausslaguerre
export gausshermite
export gaussjacobi
export gausslobatto
export gaussradau
export besselroots

include("gausslegendre.jl")
include("gausschebyshev.jl")
include("gausslaguerre.jl")
include("gausshermite.jl")
include("gaussjacobi.jl")
include("gausslobatto.jl")
include("gaussradau.jl")
include("besselroots.jl")

end


using LinearAlgebra
import SpecialFunctions: gamma, besselk

# General
abstract type Kernel end

function kernel(k::Kernel, x::AbstractVector, y::AbstractVector) end
kernel(k::Kernel, x::Number, y::Number) = kernel(k, [x], [y])

function get_params(k::Kernel) end
function set_params!(k::Kernel, params) end

# partial evaluations
kernel(k::Kernel, x::AbstractVector) = begin
    y -> kernel(k, x, y)
end

kernel(k::Kernel) = begin
    (x, y) -> kernel(k, x, y)
end

### Gaussian Radial Basis Function ###
mutable struct GaussianRBF <: Kernel
    gamma
end

get_params(k::GaussianRBF, gamma) = [gamma]
set_params!(k::GaussianRBF, gamma::Number) = begin
    k.gamma = gamma
end
set_params!(k::GaussianRBF, gamma::AbstractArray) = set_params!(k, first(gamma))

kernel(k::GaussianRBF, x::AbstractVector, y::AbstractVector) = begin
    # added factor of 0.5 to get same result as code attached to paper
    exp(- 0.5 * k.gamma^(-2) * sum((x - y).^2))
end

kernel(k::GaussianRBF, x::AbstractVector, y::AbstractVector, γ::Number) = begin
    # added factor of 0.5 to get same result as code attached to paper
    exp(- 0.5 * γ^(-2) * sum((x - y).^2))
end

### Exponential kernel
struct ExponentialKernel <: Kernel end
kernel(k::ExponentialKernel, x::AbstractVector, y::AbstractVector) = begin
    exp(dot(x, y))
end
set_params!(k::ExponentialKernel, anything) = ()

### Matern Kernel
# ### Matern Kernel
# # Works with FSSDrand and FSSDopt for SOME values, and only using `finite-differences` to optimize.
# mutable struct MaternKernel <: Kernel
#     ν
#     ρ
# end

# @inline function kernel(κ::MaternKernel, x::AbstractVector, y::AbstractVector)
#     # squared euclidean distance
#     d = norm(x - y)
#     # d = d < eps(T) ? eps(T) : d  # If d is zero, besselk will return NaN
#     d = d < eps(Float64) ? eps(Float64) : d
#     tmp = √(2κ.ν)*d /κ.ρ
#     return (2^(1.0 - κ.ν)) * (tmp^κ.ν) * besselk(κ.ν, tmp) / gamma(κ.ν)
# end
# get_params(k::MaternKernel) = [k.ν, k.ρ]
# set_params!(k::MaternKernel, ν, ρ) = begin
#     k.ν = ν
#     k.ρ = ρ
# end

# Works with FSSDrand and FSSDopt for SOME values, and only using `finite-differences` to optimize.
"""
    Matern25Kernel(ρ)

Matern25Kernel is a Matern kernel with ν = 2.5, and is thus 1st order differentiable, which is what we require.

# Notes
Optimization wrt. ρ is potentially unstable, and so some runs might return an error, while others wont.
"""
mutable struct Matern25Kernel <: Kernel
    ρ
end

@inline function kernel(κ::Matern25Kernel, x::AbstractVector, y::AbstractVector)
    ν = 2.5

    # squared euclidean distance
    d = norm(x - y)
    # d = d < eps(T) ? eps(T) : d  # If d is zero, besselk will return NaN
    d = d < eps(Float64) ? eps(Float64) : d
    tmp = √(2ν)*d /κ.ρ
    return (2^(1.0 - ν)) * (tmp^ν) * besselk(ν, tmp) / gamma(ν)
end
get_params(k::Matern25Kernel) = [k.ρ]
set_params!(k::Matern25Kernel, ρ) = begin
    k.ρ = ρ
end

### InverseMultiQuadratic (IMQ) kernel
mutable struct InverseMultiQuadratic <: Kernel
    c  # > 0
    b  # < 0
end

function kernel(k::InverseMultiQuadratic, x::AbstractVector, y::AbstractVector)
    (k.c^2 + norm(x - y)^2)^k.b
end
get_params(k::InverseMultiQuadratic) = [k.c, k.b]
set_params!(k::InverseMultiQuadratic, c, b) = begin
    k.c = c
    k.b = b
end

### Derivatives
# These functions can be implemented manually for particular kernels were analytic expressions are easily attainble,
# The below allow the use to implement any arbitrary kernel, and then we can optimize over it.

@inline k_dx(k::Kernel, x::AbstractArray, y::AbstractArray) = ForwardDiff.gradient(z -> kernel(k, z, y), x);
@inline k_dy(k::Kernel, x::AbstractArray, y::AbstractArray) = ForwardDiff.gradient(z -> kernel(k, x, z), y);
@inline k_dxdy(k::Kernel, x::AbstractArray, y::AbstractArray) = ForwardDiff.jacobian(z -> k_dx(k, x, z), y);

@inline k_dx(k::Kernel, x::Number, y::Number) = ForwardDiff.derivative(z -> kernel(k, z, y), x);
@inline k_dy(k::Kernel, x::Number, y::Number) = ForwardDiff.derivative(z -> kernel(k, x, z), y);
@inline k_dxdy(k::Kernel, x::Number, y::Number) = ForwardDiff.derivative(z -> k_dx(k, x, z), y);


"""
    VariationalDenseBE(in, out, rank,
                        ensemble_size,
                        σ=identity;
                        bias=true,
                        init=glorot_normal,
                        alpha_init=glorot_normal,
                        gamma_init=glorot_normal)
    VariationalDenseBE(layer, alpha_sampler, gamma_sampler,
                        ensemble_bias, ensemble_act, rank)

Creates a bayesian dense BatchEnsemble layer.
Batch ensemble is a memory efficient alternative for deep ensembles. In deep ensembles,
if the ensemble size is N, N different models are trained, making the time and memory
complexity O(N * complexity of one network).
BatchEnsemble generates weight matrices for each member in the ensemble using a
couple of rank 1 vectors R (alpha), S (gamma), RS' and multiplying the result with
weight matrix W element wise. We also call R and S as fast weights. In the bayesian
version of batch ensemble, instead of having point estimates of the fast weights, we
sample them form a trainable parameterized distribution.

Reference - https://arxiv.org/abs/2005.07186

During both training and testing, we repeat the samples along the batch dimension
N times, where N is the ensemble_size. For example, if each mini batch has 10 samples
and our ensemble size is 4, then the actual input to the layer has 40 samples.
The output of the layer has 40 samples as well, and each 10 samples can be considered
as the output of an esnemble member.

# Fields
- `layer`: The dense layer which transforms the pertubed input to output
- `alpha_sampler`: Sampler for the first Fast weight of size (in_dim, ensemble_size)
- `gamma_sampler`: Sampler for the second Fast weight of size (out_dim, ensemble_size)
- `ensemble_bias`: Bias added to the ensemble output, separate from dense layer bias
- `ensemble_act`: The activation function to be applied on ensemble output
- `rank`: Rank of the fast weights (rank > 1 doesn't work on GPU for now)

# Arguments
- `in::Integer`: Input dimension of features
- `out::Integer`: Output dimension of features
- `rank::Integer`: Rank of the fast weights
- `ensemble_size::Integer`: Number of models in the ensemble
- `σ::F=identity`: Activation of the dense layer, defaults to identity
- `init=glorot_normal`: Initialization function, defaults to glorot_normal
- `bias::Bool=true`: Toggle the usage of bias in the dense layer
- `ensemble_bias::Bool=true`: Toggle the usage of ensemble bias
- `ensemble_act::F=identity`: Activation function for enseble outputs
- `alpha_init=TrainableMvNormal`: Initialization function for the alpha fast weight,
                        defaults to TrainableMvNormal
- `gamma_init=TrainableMvNormal`: Initialization function for the gamma fast weight,
                        defaults to TrainableMvNormal
"""
struct VariationalDenseBE{L,F,M,B}
    layer::L
    alpha_sampler::M
    gamma_sampler::M
    ensemble_bias::B
    ensemble_act::F
    rank
end

function VariationalDenseBE(
    in::Integer,
    out::Integer,
    rank::Integer,
    ensemble_size::Integer,
    σ = identity;
    init = glorot_normal,
    bias = true,
    ensemble_bias = true,
    ensemble_act = identity,
    alpha_init = TrainableMvNormal,
    gamma_init = TrainableMvNormal,
    device = cpu
)

    layer = Flux.Dense(in, out, σ; init = init, bias = bias)
    if rank >= 1
        alpha_shape = (in, ensemble_size, rank)
        gamma_shape = (out, ensemble_size, rank)
    else
        error("Rank must be >= 1.")
    end
    alpha_sampler = alpha_init(alpha_shape, init = init, device = device)
    gamma_sampler = gamma_init(gamma_shape, init = init, device = device)
    gamma = gamma_sampler()
    ensemble_bias = create_bias(gamma, ensemble_bias, out, ensemble_size)

    return VariationalDenseBE(
        layer,
        alpha_sampler,
        gamma_sampler,
        ensemble_bias,
        ensemble_act,
        rank,
    )
end

@functor VariationalDenseBE

"""
The forward pass for a VariationalDenseBE layer. The input is initially perturbed
using the first fast weight, then passed through the dense layer, and finally
multiplied by the second fast weight.

# Arguments
- `x`: Input tensors
"""
function (be::VariationalDenseBE)(x)
    layer = be.layer
    alpha = be.alpha_sampler()
    gamma = be.gamma_sampler()
    e_b = be.ensemble_bias
    e_σ = be.ensemble_act
    rank = be.rank

    batch_size = size(x)[end]
    in_size = size(alpha)[1]
    out_size = size(gamma)[1]
    ensemble_size = size(alpha)[2]
    samples_per_model = batch_size ÷ ensemble_size

    # TODO: Do we really need repeat? Or can we use broadcast?
    # Alpha, gamma shapes - [units, ensembles, rank]
    alpha = reshape(alpha, (in_size, ensemble_size * rank))
    gamma = reshape(gamma, (out_size, ensemble_size * rank))
    # Repeat breaks on GPU when input dims > 2
    alpha = repeat(alpha, samples_per_model)
    gamma = repeat(gamma, samples_per_model)
    # Reshape alpha, gamma to [units, batch_size, rank]
    alpha = reshape(alpha, (in_size, batch_size, rank))
    gamma = reshape(gamma, (out_size, batch_size, rank))
    # Reshape inputs to [units, batch_size, 1] for broadcasting
    x = Flux.unsqueeze(x, (ndims(x) + 1))
    # Perturb the inputs
    perturbed_x = x .* alpha
    # Dense layer forward pass
    outputs = layer(perturbed_x) .* gamma
    # Reduce the rank dimension through summing it up
    outputs = sum(outputs, dims = ndims(outputs))
    outputs = reshape(outputs, (out_size, samples_per_model, ensemble_size))
    # Reshape ensemble bias
    e_b = Flux.unsqueeze(e_b, ndims(e_b))

    outputs = e_σ.(outputs .+ e_b)
    outputs = reshape(outputs, (out_size, batch_size))
    return outputs
end


### Binning
# A bin is a range of values between which values are grouped. Bins range from
# 1 to however many bins are selected. This allows for continuous data to be used
# for classification or histograms or other analysis.

using DataFrames

SampleSpaceTooSmall = AssertionError("equalfrequencylimits(dv, bins) -> bins is not less than the length of dv. Reduce bins.")
BinWidthTooSmall = AssertionError("equalfrequencylimits(dv, bins) -> The calculated width of the bins is too small for sample space. Add more data or reduce bins.")

function equaldistancelimits(dv::Vector{<:AbstractFloat}, bins::Signed)
    min = minimum(dv); max = maximum(dv); binwidth = (max-min)/bins
    return(min:binwidth:max)
end

function equalfrequencylimits(dv::Vector{<:AbstractFloat}, bins::Signed)
    if bins > length(dv)
        throw(SampleSpaceTooSmall)
    end
    sorteddv = sort(dv)
    elements = length(dv); binwidth = div(elements, bins)
    if binwidth < 1
        throw(BinWidthTooSmall)
    end
    binrange = zeros(Float64, bins)
    for i in 1:bins
        binrange[i] = dv[i*binwidth]
    end
    if mod(elements,bins) > 0 binrange[end] = dv[end] end
    return binrange
end

function find_bin(x::AbstractFloat, limits::Vector{<:AbstractFloat})
    bin = length(limits) + 1
    for i in 1:length(limits)
        if x < limits[i]
            bin = i
            break
        end
    end
    bin
end

function find_bin(x::AbstractFloat, limits::StepRangeLen{<:AbstractFloat})
    bin = length(limits) + 1
    for i in 1:length(limits)
        if x < limits[i]
            bin = i
            break
        end
    end
    bin
end

function find_bins(dv::Vector{<:AbstractFloat}, limits::Vector{<:AbstractFloat})
    vbins = zeros(Int64, length(dv))
    for i in 1:length(dv)
        vbins[i] = find_bin(dv[i], limits)
    end
    vbins
end

function find_bins(dv::Vector{<:AbstractFloat}, limits::StepRangeLen{<:AbstractFloat})
    vbins = zeros(Int64, length(dv))
    for i in 1:length(dv)
        vbins[i] = find_bin(dv[i], limits)
    end
    vbins
end

function group_bins{T}(df::AbstractDataFrame, col::T, limits::Vector{AbstractFloat})
    n_bins = length(limits) + 1

    vbins = find_bins(df[col], limits)
    (idx, starts) = DataArrays.groupsort_indexer(vbins, n_bins)

    # Remove zero-length groupings
    starts = _uniqueofsorted(starts)
    ends = [starts[2:end] - 1]
    GroupedDataFrame(df, [col], idx, starts[1:end-1], ends)
end

function group_bins{T}(df::AbstractDataFrame, col::T, limits::StepRangeLen{AbstractFloat})
    n_bins = length(limits) + 1

    vbins = find_bins(df[col], limits)
    (idx, starts) = DataArrays.groupsort_indexer(vbins, n_bins)

    # Remove zero-length groupings
    starts = unique(starts)
    ends = [starts[2:end] - 1]
    GroupedDataFrame(df, [col], idx, starts[1:end-1], ends)
end

x = [0.62, 0.99, 0.52, 0.03, 0.55, 0.74, 0.49, 0.31, 0.62, 0.19]
x_df = DataFrame(X=x, Y=rand(10))

group_bins(x_df, :X, range(0.0, 0.5, 20))


# # Optimiziation of an excitatory/inhibitory (E/I) recurrent network

#=
In this example, I optimize an RNN network with E/I units.
The dynamics is expressed as follows:
```math
\tau\, \frac{\text{d} \mathbf{u}}{\text{d}t} = - \mathbf{u} +
W f\left(\mathbf{u}\right) + \mathbf{h}
```
Where $W$ is a connection matrix with no autapses that preserves the separation
between E and I units (Dale's law). There is also a sparseness parameter for $W$,
to increase the difficulty (and as proof of concept).
Then there is a leak term $-\mathbf{u}$
and a constant external input $\mathbf{h}$. Finally the activation function
$f\left( \cdot \right)$ is a rectified-linear function.
=#

# ## Initialization
using Plots,NamedColors
using LinearAlgebra,Statistics
using Random
using SmoothedSpectralAbscissa ; const SSA=SmoothedSpectralAbscissa
using OrdinaryDiffEq # to run the dynamics
Random.seed!(0)

iofunction(x::Real) = max(0.0,x)

function run_rnn_dynamics(u0::Vector{R},W::Matrix{R},h::Vector{R},
    tmax::R,dt::R=0.01;verbose=false) where R
  f = function (du,u,p,t)
    mul!(du,W,iofunction.(u))
    @. du = du - u + h
    return du
  end
  prob = ODEProblem(f,u0,(0.,tmax))
  solv = solve(prob,Tsit5();verbose=verbose,saveat=dt)
  ret_u =hcat(solv.u...)
  ret_norms = mapslices(norm,ret_u;dims=1)[:]
  return solv.t,ret_u,ret_norms
end;

# ## Build the starting weight matrix

const sparseness = 0.5
const ne = 35
const ni = 15
const ntot = ne+ni
Wmask = hcat(fill(1.,ntot,ne),fill(-1.,ntot,ni)) # 1 for E , -1 for I, 0. for no connection
for i in 1:ntot,j in 1:ntot
  if i==j || rand()<sparseness
    Wmask[i,j]=0.
  end
end
heatmap(Wmask;
  ratio=1,seriescolor=:bwr,
  clims=(-1,1),cbar=nothing,axis=nothing,ticks=nothing,border=:none)
#=
The matrix `Wmask` shown above specifies the connectivy
(presence of a conneciton, and wether the neuron is E or I)
The full weight matrix $W$ prior to optimization is defined as follows:
=#
W0 = 2.0 .* rand(ntot,ntot) .* Wmask;
#=
Even with the leaky term, the dynamics $\mathbf{u}(t)$ is very unstable.
=#
u0 = 3.0.*randn(ntot)
h = 0.1 .* rand(ntot)
times,dyn_t,dyn_norms = run_rnn_dynamics(u0,W0,h,10.0,0.05)
plot(times,dyn_norms;
  leg=false,linewidth=3,color=:black,xlabel="time",ylabel="norm(u(t))",
  yscale=:log10,label="before optimization")

# Note that the y-scale is exponential here.

# ## Objective function that excludes diagonal
using Optim
#=
To keep the signs of $W_{i,j}$ I make a reparametrization as follows.
```math
W_{i,j} = s_j \; \exp\left(\beta_{i,j}\right)
```
I then need to propagate the gradient of the SSA. I will also inclide a
2-norm regularizer on the weights
=#
function objective_and_grad_constraints(x::Vector{R},grad::Union{Nothing,Vector{R}},
    n::Integer,ssa_eps::R,alloc::SSA.SSAAlloc,W0::Matrix{R}) where R
  betas=reshape(x,(n,n))
  mat = @. exp(betas) * W0  # transform, apply constraints
  gradmat = isnothing(grad) ? nothing : similar(mat)
  obj = SSA.ssa!(mat,gradmat,alloc,ssa_eps)
  if !isnothing(grad)
    gradmat .*= mat # propagate gradient for constraint
    for i in eachindex(gradmat)
      grad[i]=gradmat[i]
    end
  end
  return obj
end;

# ## Optimizer and optimization

const ssa_eps=0.001
const alloc = SSA.SSAAlloc(ntot)
const y0 = map(w-> w!=0. ? log(abs(w)) : 0. ,W0[:])

function objfun!(F,G,y)
  λ = 3.0/length(y) # regularizer calibration
  obj=objective_and_grad_constraints(y,G,ntot,ssa_eps,alloc,W0) # add the regularizer
  obj += 0.5*λ*mapreduce(x->x^2,+,y)
  if !isnothing(G)
    @. G += λ*y # gradient of regularizer
  end
  return obj
end

opt_out = optimize(Optim.only_fg!(objfun!),y0,BFGS(),Optim.Options(iterations=1_000))
y_opt=Optim.minimizer(opt_out)
W_opt = Wmask .* exp.(reshape(y_opt,(ntot,ntot))); # convert from beta to weight

# Comparison between inital matrix and the optimized version.
# Note the E/I separation and the spontaneous symmetry.
_ = let wboth = hcat(W0,fill(NaN,ntot,10),W_opt)
  _wex = extrema(W_opt)
  _zero_rel =abs(_wex[1])/(_wex[2]-_wex[1])
  _cgrad = cgrad([:blue,:white,:red],[0, _zero_rel ,1.])
  heatmap(wboth;ratio=1,seriescolor=_cgrad,clim=_wex,
    axis=nothing,ticks=nothing,border=:none)
end

# Now, I consider the norm of $\mathbf u(t)$ in the optimized system...
times,dyn_t_opt,dyn_norms_opt = run_rnn_dynamics(u0,W_opt,h,40.0,0.05)
plot(times,dyn_norms_opt;
      leg=:topright,linewidth=3,color=:blue,
      xlabel="time",ylabel="norm(x(t))", label="after optimization")

#=
STABLE !

There seems to be a large initial amplificaiton, but the activity settles to a stable point.
=#

# ## Extras

#=
In gradient based optimization with no automatic differentiation, it is
always necessary to test the gradient of the objective function.
The procedure is illustrated below.
=#

using Calculus
function test_gradient(myobjfun,y0)
  grad_an = similar(y0)
  _ = myobjfun(1.0,grad_an,y0) # compute gradient analytically
  grad_num = Calculus.gradient(y->myobjfun(1.0,nothing,y),y0) # compute it numerically
  for (k,y) in enumerate(y0)
    if y == 0.
      grad_num[k] = 0.
   end end
  return (grad_an,grad_num)
end

#=
To enable testing, set `do_the_test=true`
It is advised to reduce the size of the system.
=#

(x1,x2)=test_gradient(objfun!,randn(ntot^2))
scatter(x1,x2;ratio=1)
plot!(identity)

# Literate.markdown("examples/03_EI.jl","docs/src";documenter=true,repo_root_url="https://github.com/dylanfesta/SmoothedSpectralAbscissa.jl/blob/master") #src


@testset "Testing interp.jl" begin

    # uniform interpolation
    breaks = range(-3, stop=3, length=100)
    vals = exp.(breaks)
    vals2 = [vals sin.(breaks)]

    li = interp(breaks, vals)
    li2 = LinInterp(breaks, vals)

    li_mat = interp(breaks, vals2)
    li_mat2 = LinInterp(breaks, vals2)

    # test constructor
    @test li == li2
    @test li_mat == li_mat2

    # make sure evaluation is inferrable
    for T in (Float64, Float32, Float16, Int64, Int32, Int16)
        @test begin
            @inferred li(one(T))
            true
        end
        @test begin
            @inferred li_mat(one(T))
            true
        end
    end

    # on grid is exact
    for i in 1:length(breaks)
        @test abs(li(breaks[i]) - vals[i]) < 1e-15
        @test all(abs.(li_mat(breaks[i]) - vals2[i, :] .< 1e-15))
    end

    # off grid is close
    for x in range(-3, stop=3, length=300)
        @test abs(li(x) - exp(x)) < 1e-2
        @test all(abs.(li_mat(x) .- [exp(x), sin(x)]) .< 1e-2)
        @test li(x) ≈ li_mat(x, 1)
    end

    # test errors for col spec for li_mat being wrong
    @test_throws BoundsError li_mat(0.5, 0)
    @test_throws BoundsError li_mat(0.5, 3)
    @test_throws BoundsError li_mat(0.5, [0, 1])
    @test_throws BoundsError li_mat(0.5, [2, 3])

    # non-uniform
    breaks = cumsum(0.1 .* rand(20))
    vals = 0.1 .* map(sin, breaks)
    li = interp(breaks, vals)
    li_mat = interp(breaks, [vals vals .+ 1])

    # on grid is exact
    for i in 1:length(breaks)
        @test abs(li(breaks[i]) - vals[i]) < 1e-15
        @test all(abs.(li_mat(breaks[i]) .- [vals[i], vals[i]+1]) .< 1e-15)
    end

    # off grid is close
    for x in range(minimum(breaks), stop = maximum(breaks), length = 30)
        @test abs(li(x) - 0.1*sin(x)) < 1e-2
        @test all(abs.(li_mat(x) - [0.1*sin(x), 0.1*sin(x)+1]) .< 1e-2)

    end

    # un-sorted works for `interp` function, but not `LinInterp`
    breaks = rand(10)
    vals = map(sin, breaks)

    @inferred interp(breaks, vals)
    @inferred interp(breaks, [vals vals .+ 1])
    @test_throws ArgumentError LinInterp(breaks, vals)
    @test_throws ArgumentError LinInterp(breaks, [vals vals .+ 1])

    # dimension mismatch
    breaks = cumsum(rand(10))
    vals = rand(8)

    @test_throws DimensionMismatch interp(breaks, vals)
    @test_throws DimensionMismatch LinInterp(breaks, vals)

    @test_throws DimensionMismatch interp(breaks, [vals vals .+ 1])
    @test_throws DimensionMismatch LinInterp(breaks, [vals vals .+ 1])

end  # @testset


mutable struct Expo <: AbstractTweenEquation
    name::String
    # A function that accepts a time argument and returns a float
    compute::Function # f(t::Float64) -> Float64
    type::Int64

    function Expo_EaseIn(type::Int64)
        o = new()

        if type == EASE_IN
            o.type = EASE_IN
            o.name = "Expo.easeIn"
            o.compute = _exp_ease_in
        elseif type == EASE_OUT
            o.type = EASE_OUT
            o.name = "Expo.easeOut"
            o.compute = _exp_ease_out
        else
            o.type = EASE_INOUT
            o.name = "Expo.easeInOut"
            o.compute = _exp_ease_inout
        end

        o
    end
end

function _exp_ease_in(t::Float64)
    if t == 0.0
        0.0
    else
        2^(10.0 * (t - 1.0))
    end
end

function _exp_ease_out(t::Float64)
    if t == 1.0
        1.0
    else
        -(2.0^(-10.0 * t)) + 1.0
    end
end

function _exp_ease_inout(t::Float64)
    if t == 0.0
        return 0.0
    end

    if t == 1.0
        return 1.0
    end

    a = t * 2.0
    if a < 1.0
        return 0.5 * (2.0^(10.0 * (a - 1.0)))
    end

    0.5 * (-(2.0^(-10.0 * (t - 1.0))) + 2.0)
end


@inline _linear_endpoints(z::Real) = (-z, z)
@inline _linear_endpoints(z::AbstractInterval) = endpoints(z)

@inline _radial_endpoints(r::T) where {T <: Real} = (zero(T), r)
@inline _radial_endpoints(r::AbstractInterval) = endpoints(r)

@inline _in_angular_interval_closed(α::Real, α_int::Nothing, tol::Real = 0) = true
@inline _in_angular_interval_closed(α::Real, α_int::AbstractInterval{T}, tol::Real = 0) where {T} = mod(α - (α_int.left-tol), T(2π)) ≤ (α_int.right+tol) - (α_int.left-tol)

@inline function _in_angular_interval_closed(α::T, α_int::Tuple{T,T}; csgtol::T = csg_default_tol(T)) where {T}
    m = mod(α - α_int[1], T(2π))
    d = (α_int[2] - α_int[1])
    m ≤ d + csgtol
end
@inline function _in_angular_interval_open(α::T, α_int::Tuple{T,T}; csgtol::T = csg_default_tol(T)) where {T}
    csgtol < mod(α - α_int[1], T(2π)) < (α_int[2] - α_int[1]) - csgtol
end

@inline _in_angular_interval_open(α::T, α_int::Nothing) where {T<:Real} = 0 < mod(α, T(2π)) < T(2π)
@inline _in_angular_interval_open(α::Real, α_int::AbstractInterval{T}) where {T} = 0 < mod(α - α_int.left, T(2π)) < width(α_int)


function dstats(x, rnames="";short=false)
    k = size(x,2)
    if rnames==""
        rnames = 1:k
        rnames = rnames'
    end
    m = mean(x,dims=1)
    mm = median(x,dims=1)
    s = std(x,dims=1)
    sk = m-m
    kt = m-m
    mn = minimum(x,dims=1)
    mx = maximum(x,dims=1)
    q05 = fill(0.0,k)
    q95 = fill(0.0,k)
    if short == false
        for i = 1:size(x,2) q05[i] = quantile(x[:,i], 0.05) end
        for i = 1:size(x,2) q95[i] = quantile(x[:,i], 0.95) end
        cnames = ["  mean", " median","  std", "min", "max", "q05", "q95"]
        stats = [m' mm' s' mn' mx' q05 q95]
        prettyprint(stats, cnames, rnames)
    else
        cnames = ["  mean", " median", "  std", "min", "max"]
        stats = [m' mm' s' mn' mx']
        prettyprint(stats, cnames, rnames)
    end
    return stats
end


using LinearAlgebra
using QuasiTriangular
using Random
using Test

Aorig = [1 3; 0 2]
A1 = [Aorig Aorig; zeros(2,2) Aorig]
A = QuasiUpperTriangular(A1)
B = randn(4,4)
C = zeros(4,4)
#A_mul_B!(C, 1, B, 1, 4, 4, A, 1, 4)

#@test C ≈ B*A1

Aorig = [1 3; 0.5 2]
A = QuasiUpperTriangular(Aorig)
B = Matrix{Float64}(I(2))
X = zeros(2,2)
#A_ldiv_B!(A,B)

#@test B ≈ inv(Aorig)

B = Matrix{Float64}(I(2))
A_rdiv_Bt!(B,A)
@test B ≈ inv(Aorig)'

B = Matrix{Float64}(I(2))
A_rdiv_B!(B,A)

@test B ≈ inv(Aorig)

Random.seed!(123)
n = 7
a = randn(n,n)
S = schur(a)
t = S.T
b = randn(n,n)
c = similar(b)

@test t*b ≈ A_mul_B!(QuasiUpperTriangular(t),b)
@test t'*b ≈ At_mul_B!(QuasiUpperTriangular(t),b)
@test b*t ≈ A_mul_B!(b,QuasiUpperTriangular(t))
@test b*t' ≈ A_mul_Bt!(b,QuasiUpperTriangular(t))

@test t*b ≈ A_mul_B!(c,QuasiUpperTriangular(t),b)
@test t'*b ≈ At_mul_B!(c,QuasiUpperTriangular(t),b)
@test b*t ≈ A_mul_B!(c,b,QuasiUpperTriangular(t))
@test b*t' ≈ A_mul_Bt!(c,b,QuasiUpperTriangular(t))

b1 = copy(b)
x = zeros(n,n)
A_ldiv_B!(QuasiUpperTriangular(t),b1)
@test t\b ≈ b1
b1 = copy(b)
A_rdiv_B!(b1,QuasiUpperTriangular(t))
@test b/t ≈ b1
b1 = copy(b)
A_rdiv_Bt!(b1,QuasiUpperTriangular(t))
@test b/t' ≈ b1

b = rand(n)
b1 = copy(b)
r = rand()
I_plus_rA_ldiv_B!(r,QuasiUpperTriangular(t),b1)
@test b1 ≈ (Matrix{Float64}(I(n)) + r*t)\b
b1 = copy(b)
s = rand()
I_plus_rA_plus_sB_ldiv_C!(r,s,QuasiUpperTriangular(t),QuasiUpperTriangular(t*t),b1)
@test b1 ≈ (Matrix{Float64}(I(n)) + r*t + s*t*t)\b

b = rand(n,n)
b1 = copy(b)
r = rand()
I_plus_rA_ldiv_B!(r,QuasiUpperTriangular(t),b1)
@test b1 ≈ (Matrix{Float64}(I(n)) + r*t)\b
b1 = copy(b)
s = rand()
I_plus_rA_plus_sB_ldiv_C!(r,s,QuasiUpperTriangular(t),QuasiUpperTriangular(t*t),b1)
@test b1 ≈ (Matrix{Float64}(I(n)) + r*t + s*t*t)\b


type QPTrace
    x
    xc
    value
    search
    clamped
    nfactor
end

"""
 Minimize `0.5*x'*H*x + x'*g`  s.t. lower<=x<=upper

  inputs:
     `H`            - positive definite matrix   (n * n)
     `g`            - bias vector                (n)
     `lower`        - lower bounds               (n)
     `upper`        - upper bounds               (n)

   optional inputs:
     `x0`           - initial state              (n)
     `options`      - see below                  (7)

  outputs:
     `x`            - solution                   (n)
     `result`       - result type (roughly, higher is better, see below)
     `Hfree`        - subspace cholesky factor   (n_free * n_free)
     `free`         - set of free dimensions     (n)
"""
function boxQP(H,g,lower,upper,x0;
               maxIter        = 100,
               minGrad        = 1e-8,
               minRelImprove  = 1e-8,
               stepDec        = 0.6,
               minStep        = 1e-22,
               Armijo         = 0.1,
               print          = 0)
    #     maxIter        = 100        # maximum number of iterations
    #     minGrad        = 1e-8       # minimum norm of non-fixed gradient
    #     minRelImprove  = 1e-8       # minimum relative improvement
    #     stepDec        = 0.6        # factor for decreasing stepsize
    #     minStep        = 1e-22      # minimal stepsize for linesearch
    #     Armijo         = 0.1    	# Armijo parameter (fraction of linear improvement required)
    #     print          = 0 			# verbosity

    n        = size(H,1)
    clamped  = falses(n)
    free     = trues(n)
    oldvalue = 0.
    result   = 0
    gnorm    = 0.
    nfactor  = 0
    trace    = Array(QPTrace,maxIter)
    Hfree    = zeros(n,n)
    clamp(x) = max(lower, min(upper, x))

    debug("# initial state")
    x = clamp(x0[:])
    LU = [lower upper]
    LU[!isfinite(LU)] = NaN

    debug("# initial objective value")
    value    = (x'*g + 0.5*x'*H*x )[1]

    if print > 0
        @printf("==========\nStarting box-QP, dimension %-3d, initial value: %-12.3f\n",n, value)
    end

    debug("# main loop")
    iter = 1
    while iter <= maxIter

        if result != 0
            break
        end

        debug("# check relative improvement")
        if iter>1 && (oldvalue - value) < minRelImprove*abs(oldvalue)
            result = 4
            break
        end
        oldvalue = value

        debug("# get gradient")
        grad     = g + H*x

        debug("# find clamped dimensions")
        old_clamped                        = clamped
        clamped                            = falses(n)
        #         clamped[(x[:,1] .== lower)&(grad[:,1].>0)]   = true
        #         clamped[(x[:,1] .== upper)&(grad[:,1].<0)]   = true
        for i = 1:n
            clamped[i]   = ((x[i,1] == lower[i])&&(grad[i,1]>0)) || ((x[i,1] == upper[i])&&(grad[i,1]<0))
        end
        free                               = !clamped

        debug("# check for all clamped")
        if all(clamped)
            result = 6
            break
        end

        debug("# factorize if clamped has changed")
        if iter == 1
            factorize    = true
        else
            factorize    = any(old_clamped != clamped)
        end

        if factorize
            Hfree  = chol(H[free,free])  # was   (Hfree, indef)  = chol(H[free,free])
            #             if indef
            #                 result = -1
            #                 break
            #             end
            nfactor += 1
        end

        debug("# check gradient norm")
        gnorm  = norm(grad[free])
        if gnorm < minGrad
            result = 5
            break
        end

        debug("# get search direction")
        grad_clamped   = g  + H*(x.*clamped)
        search         = zeros(n,1)
        search[free]   = -Hfree\(Hfree'\grad_clamped[free]) - x[free]

        debug("# check for descent direction")
        sdotg          = sum(search.*grad)
        if sdotg >= 0 # (should not happen)
            break
        end

        debug("# armijo linesearch")
        step  = 1
        nstep = 0
        xc    = clamp(x+step*search)
        vc    = (xc'*g + 0.5*xc'*H*xc)[1]
        while (vc - oldvalue)/(step*sdotg) < Armijo
            step  = step*stepDec
            nstep += 1
            xc    = clamp(x+step*search)
            vc    = (xc'*g + 0.5*xc'*H*xc)[1]
            if step<minStep
                result = 2
                break
            end
        end

        if print > 1
            @printf("iter %-3d  value % -9.5g |g| %-9.3g  reduction %-9.3g  linesearch %g^%-2d  n_clamped %d\n",
                    iter, vc, gnorm, oldvalue-vc, stepDec, nstep, sum(clamped))
        end

        trace[iter] = QPTrace(x,xc,value,search,clamped,nfactor )

        debug("# accept candidate")
        x     = xc
        value = vc
        iter += 1

    end

    if iter == maxIter
        result = 1
    end

    results = ["Hessian is not positive definite",          # result = -1
               "No descent direction found",                # result = 0    SHOULD NOT OCCUR
               "Maximum main iterations exceeded",          # result = 1
               "Maximum line-search iterations exceeded",   # result = 2
               "No bounds, returning Newton point",         # result = 3
               "Improvement smaller than tolerance",        # result = 4
               "Gradient norm smaller than tolerance",      # result = 5
               "All dimensions are clamped"]                # result = 6

    if print > 0
        @printf("RESULT: %s.\niterations %d  gradient %-12.6g final value %-12.6g  factorizations %d\n",
                results[result+2], iter, gnorm, value, nfactor)
    end

    return x,result,Hfree,free,trace
end

function demoQP()
    n 		= 500
    g 		= randn(n,1)
    H 		= randn(n,n)
    H 		= H*H'
    lower 	= -ones(n,1)
    upper 	=  ones(n,1)
    @time boxQP(H, g, lower, upper, randn(n,1),print=1)

end


module KalmanFilter

#
# References:
# Haykin, Kalman Filter and NN (2001)
# https://en.wikipedia.org/wiki/Kalman_filter#Fixed-interval_smoothers
# Hartikainen et al., Optimal filtering with Kalman filters and smoothers - A Manual for Matlab toolbox EKF/UKF (2011)
# Martin, Modified Bryson-Frazier Smoother Cross-Covariance. IEEE TRANSACTIONS ON AUTOMATIC CONTROL, VOL. 59, NO. 1, JANUARY 2014
#

using LinearAlgebra

include("CommonStructures.jl")
export CholeskySqrt, SVDSqrt, CholeskyModSqrt

# Review the use of eye, because it could be replaced by I
# unpack and unpack_matrix could be merged into just unpack accepting different methods
include("Utils.jl")
export eye, unpack, unpack_matrix

include("FixedIntervalSmoothers.jl")
export kalman_smooth,
       RauchTungStriebel, ModifiedBrysonFrazier,
       RTSSmoother, MBFSmoother

include("LinearKalmanFilter.jl")
export linear_kalman, LinearKF

include("UnscentedKalmanFilter.jl")
export unscented_kalman, UnscentedKF

end # module KalmanFilter


"""
RationalQuadraticKernel([ρ=1.0[,α=2.0]])
The rational-quadratic kernel is an isotropic Mercer kernel given by the formula:
```
    κ(x,y)=(1+ρ²||x−y||²/α)^(-α)
```
where `α` is a shape parameter of the Euclidean distance. Check [`GammaRationalQuadraticKernel`](@ref) for a generalization.
"""
struct RationalQuadraticKernel{Tr,Tα<:Real} <: Kernel{Tr}
    transform::Tr
    α::Tα
    function RationalQuadraticKernel{Tr, Tα}(t::Tr, α::Tα) where {Tr, Tα}
        @check_args(RationalQuadraticKernel, α, α > zero(Tα), "α > 1")
        return new{Tr, Tα}(t, α)
    end
end

function RationalQuadraticKernel(ρ::Real=1.0, α::Real=2.0)
    RationalQuadraticKernel(ScaleTransform(ρ),α)
end

function RationalQuadraticKernel(ρ::AbstractVector{<:Real}, α::Real=2.0)
    RationalQuadraticKernel(ARDTransform(ρ), α)
end

function RationalQuadraticKernel(t::Tr, α::Tα=2.0) where {Tr<:Transform, Tα<:Real}
    return RationalQuadraticKernel{Tr, Tα}(t, α)
end

params(k::RationalQuadraticKernel) = (params(transform(k)),k.α)
opt_params(k::RationalQuadraticKernel) = (opt_params(transform(k)),k.α)

@inline kappa(κ::RationalQuadraticKernel, d²::T) where {T<:Real} = (one(T)+d²/κ.α)^(-κ.α)

metric(::RationalQuadraticKernel) = SqEuclidean()

"""
`GammaRationalQuadraticKernel([ρ=1.0[,α=2.0[,γ=2.0]]])`
The Gamma-rational-quadratic kernel is an isotropic Mercer kernel given by the formula:
```
    κ(x,y)=(1+ρ^(2γ)||x−y||^(2γ)/α)^(-α)
```
where `α` is a shape parameter of the Euclidean distance and `γ` is another shape parameter.
"""
struct GammaRationalQuadraticKernel{Tr, Tα<:Real, Tγ<:Real} <: Kernel{Tr}
    transform::Tr
    α::Tα
    γ::Tγ
    function GammaRationalQuadraticKernel{Tr,Tα,Tγ}(t::Tr, α::Tα, γ::Tγ) where {Tr, Tα<:Real, Tγ<:Real}
        @check_args(GammaRationalQuadraticKernel, α, α > one(Tα), "α > 1")
        @check_args(GammaRationalQuadraticKernel, γ, γ >= one(Tγ), "γ >= 1")
        return new{Tr, Tα, Tγ}(t, α, γ)
    end
end

function GammaRationalQuadraticKernel(ρ::Real=1.0, α::Real=2.0, γ::Real=2.0)
    GammaRationalQuadraticKernel(ScaleTransform(ρ), α, γ)
end

function GammaRationalQuadraticKernel(ρ::AbstractVector{<:Real}, α::Real=2.0, γ::Real=2.0)
    GammaRationalQuadraticKernel(ARDTransform(ρ),α,γ)
end

function GammaRationalQuadraticKernel(t::Tr,α::Tα=2.0,γ::Tγ=2.0) where {Tr<:Transform, Tα<:Real, Tγ<:Real}
    GammaRationalQuadraticKernel{Tr, Tα, Tγ}(t, α, γ)
end

params(k::GammaRationalQuadraticKernel) = (params(k.transform),k.α,k.γ)
opt_params(k::GammaRationalQuadraticKernel) = (opt_params(k.transform),k.α,k.γ)

@inline kappa(κ::GammaRationalQuadraticKernel, d²::T) where {T<:Real} = (one(T)+d²^κ.γ/κ.α)^(-κ.α)

metric(::GammaRationalQuadraticKernel) = SqEuclidean()


export ModelKitHomotopy

"""
    ModelKitHomotopy(system::ModelKit.System) <: AbstractSystem
"""
struct ModelKitHomotopy{T,P<:Union{Nothing,AbstractVector}} <: AbstractHomotopy
    homotopy::ModelKit.CompiledHomotopy{T}
    parameters::P
end

function ModelKitHomotopy(H::ModelKit.Homotopy; parameters = nothing)
    ModelKitHomotopy(ModelKit.compile(H), parameters)
end

Base.size(H::ModelKitHomotopy) = size(H.homotopy)

function set_parameters!(H::ModelKitHomotopy, p::Tuple, γ::Union{Nothing,Tuple} = nothing)
    p1, p2 = p
    n = length(p1)
    for i = 1:n
        H.parameters[i] = p1[i]
    end
    for i = 1:n
        H.parameters[i+n] = p2[i]
    end

    if γ !== nothing
        γ₁, γ₂ = γ
        H.parameters[2n+1] = γ₁
        H.parameters[2n+2] = γ₂
    end

    H
end

"""
    set_start_parameters!(H::ModelKitHomotopy, p)

Update the start parameters of `H`. This assumes that the first half of the parameters
are the start parameters.
"""
function set_start_parameters!(H::ModelKitHomotopy, p)
    for i = 1:div(length(p), 2)
        H.parameters[i] = p[i]
    end
    H
end

"""
    set_target_parameters!(H::ModelKitHomotopy, p)

Update the target parameters of `H`. This assumes that the second half of the parameters
are the start parameters.
"""
function set_target_parameters!(H::ModelKitHomotopy, p)
    n = div(length(p), 2)
    for i = 1:n
        H.parameters[i+n] = p[i]
    end
    H
end

cache(H::ModelKitHomotopy, x, t, p = nothing) = HomotopyNullCache()

function evaluate!(u, H::ModelKitHomotopy, x, t, ::HomotopyNullCache)
    ModelKit.evaluate!(u, H.homotopy, x, t, H.parameters)
end
function evaluate(H::ModelKitHomotopy, x, t, ::HomotopyNullCache)
    convert(Vector, ModelKit.evaluate(H.homotopy, x, t, H.parameters))
end
(H::ModelKitHomotopy)(x, t) = evaluate(F, x, t, H.parameters)

function jacobian!(U, H::ModelKitHomotopy, x, t, ::HomotopyNullCache)
    ModelKit.jacobian!(U, H.homotopy, x, t, H.parameters)
end
function jacobian(H::ModelKitHomotopy, x, t, ::HomotopyNullCache)
    convert(Matrix, ModelKit.jacobian(H.homotopy, x, t, H.parameters))
end

function evaluate_and_jacobian!(u, U, H::ModelKitHomotopy, x, t, ::HomotopyNullCache)
    ModelKit.evaluate_and_jacobian!(u, U, H.homotopy, x, t, H.parameters)
end

function dt!(u, H::ModelKitHomotopy, x, t, ::HomotopyNullCache)
    ModelKit.dt!(u, H.homotopy, x, t, H.parameters)
end
function dt(H::ModelKitHomotopy, x, t, ::HomotopyNullCache)
    convert(Vector, ModelKit.dt(H.homotopy, x, t, H.parameters))
end

function jacobian_and_dt!(U, u, H::ModelKitHomotopy, x, t, ::HomotopyNullCache)
    ModelKit.jacobian_and_dt!(U, u, H.homotopy, x, t, H.parameters)
end
function jacobian_and_dt(H::ModelKitHomotopy, x, t, ::HomotopyNullCache)
    ModelKit.jacobian_and_dt(H.homotopy, x, t, H.parameters)
end


import Base.length
import Base.show
import Base.size

flatten(arr) = reshape(arr, length(arr))

typealias Volume{T} Array{T, 3}

type Vol
    w::Volume{Float64}
    dw::Volume{Float64}

    function Vol(dimensions, value::Float64)
        w = Array(Float64, dimensions)
        dw = Array(Float64, dimensions)
        new(ones(dimensions) * value, zeros(dimensions))
    end

    Vol(w::Array{Float64, 3}, dw::Array{Float64, 3}) = new(w, dw)

    function Vol(dimensions)
        (x, y, depth) = dimensions
        scale = sqrt(1.0 / (x * y * depth))
        new(randn(x, y, depth) * scale, zeros(x, y, depth))
    end
end

function show(io::IO, v::Vol)
    println("Vol")
    println(io, "w: ", v.w)
    println(io, "dw: ", v.dw)
end

function length(v::Vol)
    @assert size(v.w) == size(v.dw)
    length(v.w)
end

function size(v::Vol)
    @assert size(v.w) == size(v.dw)
    size(v.w)
end