using Revise
using IMPES
using Plots

# using Plots, LinearAlgebra, IMPES
# using AnalyticalEOR

#* Define Relative Perms
kr_ww = SingleRelPerms(
        swr= 0.11,
        sor= 0.10,
        krw0= 0.2,
        kro0= 0.8,
        nw= 3.0,
        no= 2.0,
            )

kr_ow = SingleRelPerms(
        swr= 0.11,
        sor= 0.16,
        krw0= 0.4,
        kro0= 0.5,
        nw= 2.0,
        no= 3.0,
        )

kr_mw = MixedRelPerms(kr_ww, kr_ow)

#* Define reservoir
res = Reservoir(
    L = 300.,
    A = 200 * 50 * 0.3048^2,
    k = 100.0 * 9.869233e-16,
    phi = 0.2,
    )

#* Define fluids
fluid = Fluids(
            μw=1e-3,
            μo=5e-3,
)

#* Define grid
grid = Grid(
            res=res,
            Nx = 100,
            )

#* Define BC
bc = BoundaryConditions(
                        [:neumann, :dirichlet],
                        [426.5 * 0.3048 ^ 3 / 3600 / 24, 1000 * 6894.76]
                        )

PV = res.A * res.L * res.phi
tmax = PV / bc.value[1] * 3
dt = res.A * grid.dx / bc.value[1] / 50
Int(tmax ÷ dt) + 1

#* Initialize simulation and postprocessing
sim = Simulation(
                1000*ones(grid.Nx)* 6894.76,
                0.2*ones(grid.Nx),
                0.5*ones(grid.Nx),
                tmax,
                dt
)

sim_ww = deepcopy(sim)
sim_ow = deepcopy(sim)

post = FlowResults(sim, grid, tmax, dt)
post_ww = FlowResults(sim, grid, tmax, dt)
post_ow = FlowResults(sim, grid, tmax, dt)

while sim.t <= tmax
    IMPES!(sim, res, kr_mw, fluid, grid, bc, post)
    IMPES!(sim_ww, res, kr_ww, fluid, grid, bc, post_ww)
    IMPES!(sim_ow, res, kr_ow, fluid, grid, bc, post_ow)
end

i = 200
plot(grid.x, post.s[:,i], seriestype=:scatter, label="Mixed-Wet", alpha=0.5)
plot!(grid.x, post_ww.s[:,i], seriestype=:scatter, label="Water-Wet", alpha=0.5)
plot!(grid.x, post_ow.s[:,i], seriestype=:scatter, label="Oil-Wet", alpha=0.5)

#* Recovery Factors
RF = recoveryFactor(grid, post)
RF_ww = recoveryFactor(grid, post_ww)
RF_ow = recoveryFactor(grid, post_ow)
PVs = injPoreVolumes(bc, res, sim)

plot(PVs, RF, label="Mixed Wet", lw=3, alpha=0.5)
plot!(PVs, RF_ww, label="Water Wet", lw=3, alpha=0.5)
plot!(PVs, RF_ow, label="Oil Wet", lw=3, alpha=0.5)

examples/Ex2 - BL1D Mixed Wettability.jl


module AverageShiftedHistograms

import StatsBase
import StatsBase: nobs, fit!
import Distributions
import UnicodePlots
using RecipesBase

export
    ash, fit!, extendrange, xy, xyz, nout, nobs,
    # Kernels
    Kernels,
    biweight, cosine, epanechnikov, triangular, tricube, triweight, uniform,
    gaussian, logistic

# common
abstract AbstractAsh
StatsBase.nobs(o::AbstractAsh) = o.nobs
nout(o::AbstractAsh) = nobs(o) - sum(o.v)

"""
`extendrange(x, s = .5, n = 200)`

Create a `LinSpace` of length `n` starting at `s` standard deviations below
`minimum(x)` and ending at `s` standard deviations above `maximum(x)`
"""
function extendrange(y::AbstractVector, s = 0.5, n = 200)
    σ = std(y)
    linspace(minimum(y) - s * σ, maximum(y) + s * σ, n)
end

include("kernels.jl")
include("univariate.jl")
include("bivariate.jl")
end


using DataFrames, MANOVABNPTest, StatsBase, LinearAlgebra, StaticArrays, Random
rng = MersenneTwister(1)

function test_sample_01(D::Int, h::Int, l::Int, H0::Int)
    N = 4 * [50, 150, 300]
    K = [
        0.625 1.498 1.162;
        0.819 1.806 1.359;
        1.051 1.994 1.652
    ]
    γ = γvector(4, H0)
    Σ = cholesky(collect(I(D)))
    x = 1 .+ (0:N[h]-1) .% 4
    y = rand(N[h], D) * Σ.U
    for i in 1:N[h]
        x[i] == 2 && (y[i, :] .+= K[l, 1] * γ[2])
        x[i] == 3 && (y[i, :] .*= K[l, 2] ^ γ[3])
        x[i] == 4 && (y[i, :] .+= K[l, 3] * γ[4] * rand([-1, 1]))
    end
    ȳ, S = mean_and_cov(y)
    y = (y .- ȳ) / cholesky(S).U
    y = [SVector{D}(y[i, :]) for i ∈ 1:N[h]]
    return y, x
end

function test_sample_02(D::Int, h::Int, l::Int, H0::Int)
    N = 4 * [50, 150, 300]
    K = [
        0.625 1.498 1.162;
        0.819 1.806 1.359;
        1.051 1.994 1.652
    ]
    γ = γvector(4, H0)
    Σ = cholesky(collect(I(D)))
    x = 1 .+ (0:N[h]-1) .% 4
    y = rand(N[h], D) * Σ.U
    for i in 1:N[h]
        x[i] == 2 && (y[i, :] .+= K[l, 1] * γ[2])
        x[i] == 3 && (y[i, :] .*= K[l, 2] ^ γ[3])
        x[i] == 4 && (y[i, :] .+= K[l, 3] * γ[4] * rand([-1, 1]))
    end
    return y, x
end

# Exp 1
D = 2
m = MANOVABNPTest.Model(D = D)
y, x = test_sample_01(D, 1, 1, 2)
pγ = MANOVABNPTest.fit(m, y, x; iter = 4000, refresh_rate = 50, rng = rng)
println(pγ)

# Exp 2
D = 2
m = MANOVABNPTest.Model(D = D)
y, x = test_sample_01(D, 1, 1, 3)
pγ = MANOVABNPTest.fit(m, y, x; iter = 4000, refresh_rate = 50, rng = rng)
println(pγ)

# Exp 3
D = 10
m = MANOVABNPTest.Model(D = D)
y, x = test_sample_01(D, 1, 1, 2)
pγ = MANOVABNPTest.fit(m, y, x; iter = 4000, refresh_rate = 50, rng = rng)
println(pγ)

# Exp 4
D = 10
m = MANOVABNPTest.Model(D = D)
y, x = test_sample_01(D, 1, 1, 3)
pγ = MANOVABNPTest.fit(m, y, x; iter = 4000, refresh_rate = 50, rng = rng)
println(pγ)

# Exp 5
D = 10
m = MANOVABNPTest.Model(D = D)
y, x = test_sample_01(D, 1, 2, 2)
pγ = MANOVABNPTest.fit(m, y, x; iter = 4000, refresh_rate = 50, rng = rng)
println(pγ)

# Exp 6
D = 10
m = MANOVABNPTest.Model(D = D)
y, x = test_sample_01(D, 1, 2, 3)
pγ = MANOVABNPTest.fit(m, y, x; iter = 4000, refresh_rate = 50, rng = rng)
println(pγ)

# Exp 7
D = 10
m = MANOVABNPTest.Model(D = D)
y, x = test_sample_01(D, 1, 3, 2)
pγ = MANOVABNPTest.fit(m, y, x; iter = 4000, refresh_rate = 50, rng = rng)
println(pγ)

# Exp 8
D = 10
m = MANOVABNPTest.Model(D = D)
y, x = test_sample_01(D, 1, 3, 3)
pγ = MANOVABNPTest.fit(m, y, x; iter = 4000, refresh_rate = 50, rng = rng)
println(pγ)

# grid = LinRange(-3, 3, 10) |> collect
# m = MANOVABNPTest.Model(D = 2)
# MANOVABNPTest.fit(m, y, x, grid; iter = 200, rng = rng);

# Exp 9
y, x = test_sample_01(2, 1, 1, 2)
Y = Matrix{Float64}(vcat(y'...))
println(train(Y, x))

# Exp 10
y, x = test_sample_01(2, 1, 1, 2)
Y = Matrix{Float64}(vcat(y'...))
println(train_γ(Y, x))


cla()
x = range(0, stop=π, length=250)
for α ∈ [1, 4, 7, 10]
    y = @. exp(-√α*x)*sin(2*x/α)*α^2
    plot!(x, y, ls="-", lw=0.05)
end
mx = x[1:10:end]
scatter!(mx, sin.(mx), marker="star", col="salmon", msize=0.3)
xlim(0, π)


abstract type EsatMethod end
struct Sonntag1990 <: EsatMethod end
struct Alduchov1996 <: EsatMethod end
struct Allen1998 <: EsatMethod end

"""
    Esat_slope(Tair; Esat_formula, constants)
    Esat_from_Tair(Tair; Esat_formula, constants)
    Esat_from_Tair_deriv(Tair; Esat_formula, constants)

Saturation Vapor Pressure (Esat) and Slope of the Esat Curve

# Arguemtns
- `Tair`:      Air temperature (deg C)
- `Esat_formula=Sonntag1990()`:   Esat_formula to be used. Either
   `Sonntag1990()` (Default), `Alduchov1996()`, or `Allen1998()`.
- `constants=BigleafConstants()`: Dictionary with entry  :Pa2kPa

# Details
Esat (kPa) is calculated using the Magnus equation:

```Esat = a * exp((b * Tair) / (c + Tair)) / 1000```
where the coefficients a, b, c take different values depending on the Esat_formula use
The default values are from Sonntag 1990 (a=611.2, b=17.62, c=243.12). This versi
of the Magnus equation is recommended by the WMO (WMO 2008; p1.4-29). Alternativel
parameter values determined by Alduchov & Eskridge 1996 or Allen et al. 1998 can b
used (see references).
The slope of the Esat curve (``\\Delta``) is calculated as the first derivative of the function:

``\\Delta = {dEsat \\over dTair}``

# Value
- `Esat_from_Tair`: Saturation vapor pressure (kPa)
- `Esat_from_Tair_deriv`: Slope of the saturation vapor pressure curve (kPa K-1)
- `Esat_slope`: A tuple of the two above values
# References
Sonntag D. 1990: Important new values of the physical constants of 1986, vapor
pressure formulations based on the ITS-90 and psychrometric formulae.
Zeitschrift fuer Meteorologie 70, 340-344.

World Meteorological Organization 2008: Guide to Meteorological Instruments
and Methods of Observation (WMO-No.8). World Meteorological Organization,
Geneva. 7th Edition,

Alduchov, O. A. & Eskridge, R. E., 1996: Improved Magnus form approximation of
saturation vapor pressure. Journal of Applied Meteorology, 35, 601-609

Allen, R.G., Pereira, L.S., Raes, D., Smith, M., 1998: Crop evapotranspiration -
Guidelines for computing crop water requirements - FAO irrigation and drainage
paper 56, FAO, Rome.

```@example
Esat_from_Tair(20.0)          # Esat in kPa
Esat_from_Tair_deriv(20.0)    # its derivative to temperature in kPa K-1
Esat_slope(20.0)              # both as a tuple
```
"""
function Esat_slope(Tair::Number; Esat_formula=Sonntag1990(), constants=BigleafConstants())
  Esat = Esat_from_Tair(Tair; Esat_formula, constants)
  Delta = Esat_from_Tair_deriv(Tair; Esat_formula, constants)
  Esat, Delta
end,
function Esat_from_Tair(Tair; Esat_formula=Sonntag1990(), constants=BigleafConstants())
  a,b,c = get_EsatCoef(Esat_formula)
  Esat = a * exp((b * Tair) / (c + Tair)) * constants.Pa2kPa
end,
function Esat_from_Tair_deriv(Tair; Esat_formula=Sonntag1990(), constants=BigleafConstants())
  # slope of the saturation vapor pressure curve
  #Delta = eval(D(expression(a * exp((b * Tair) / (c + Tair))),name="Tair"))
  a,b,c = get_EsatCoef(Esat_formula)
  #Delta_Pa = @. a*(b / (Tair + c) + (-Tair*b) / ((Tair + c)^2))*exp((Tair*b) / (Tair + c))
  Delta_Pa = @. a * (exp((b * Tair)/(c + Tair)) * (b/(c + Tair) - (b * Tair)/(c + Tair)^2))
  Delta = Delta_Pa .* constants.Pa2kPa
end

get_EsatCoef(::Sonntag1990) = (a=611.2,b=17.62,c=243.12)
get_EsatCoef(::Alduchov1996) = (a=610.94,b=17.625,c=243.04)
get_EsatCoef(::Allen1998) = (a=610.8,b=17.27,c=237.3)

"""
    LE_to_ET(LE,Tair)
    ET_to_LE(ET,Tair)

Convert evaporative water flux from mass (ET=evapotranspiration)
             to energy (LE=latent heat flux) units, or vice versa.

# Arguments
- LE   Latent heat flux (W m-2)
- ET   Evapotranspiration (kg m-2 s-1)
- Tair Air temperature (deg C)

# Details
The conversions are given by:
- ``ET = LE/\\lambda``
- ``LE = \\lambda ET``

where ``\\lambda`` is the latent heat of vaporization (J kg-1) as calculated by
[`latent_heat_vaporization`](@ref).

# Examples
```jldoctest; output = false
# LE of 200 Wm-2 and air temperature of 25degC
ET = LE_to_ET(200,25)
≈(ET, 8.19e-5, atol =1e-7)
# output
true
```
"""
function LE_to_ET(LE,Tair)
  lambda = latent_heat_vaporization(Tair)
  ET     = LE/lambda
end,
function ET_to_LE(ET,Tair)
  lambda = latent_heat_vaporization(Tair)
  LE     = ET*lambda
end

"""
    ms_to_mol(G_ms,Tair,pressure; constants=BigleafConstants())
    mol_to_ms(G_mol,Tair,pressure; constants=BigleafConstants())

Converts conductances from mass (m s-1) to molar units (mol m-2 s-1), or vice versa

# Details
The conversions are given by
- ``G_{mol} = G_{ms}  \\, pressure / (Rgas  Tair)``
- ``G_{ms} = G_{mol}  \\, (Rgas  Tair) / pressure``

where Tair is in Kelvin and pressure in Pa (converted from kPa internally).

# References
Jones, H_G_ 1992_ Plants and microclimate: a quantitative approach to
  environmental plant physiology_
  2nd Edition_, Cambridge University Press, Cambridge_ 428

# Examples
```jldoctest; output = false
G_ms,Tair,pressure = 0.005,25.0,100.0
rmol = ms_to_mol(G_ms,Tair,pressure)
≈(rmol, 0.2017, atol =1e-4)
# output
true
```
"""
function ms_to_mol(G_ms,Tair,pressure; constants=BigleafConstants())
  Tair     = Tair + oftype(Tair,constants.Kelvin)
  pressure = pressure * constants.kPa2Pa
  G_mol  = G_ms * pressure / (oftype(Tair, constants.Rgas)  * Tair)
end,
function mol_to_ms(G_mol,Tair,pressure; constants=BigleafConstants())
  Tair     = Tair + oftype(Tair,constants.Kelvin)
  pressure = pressure * constants.kPa2Pa
  G_ms  = G_mol * (oftype(Tair, constants.Rgas) * Tair) / (pressure)
end

"""
    VPD_to_rH(VPD,Tair; ...)
    H_to_VPD(rH,Tair; ...)
    e_to_rH(e,Tair; ...)
    VPD_to_e(VPD,Tair; ...)
    e_to_VPD(e,Tair; ...)
    e_to_q(e,pressure; ...)
    q_to_e(q,pressure; ...)
    q_to_VPD(q,Tair,pressure; ...)
    VPD_to_q(VPD,Tair,pressure; ...)

Conversion between vapor pressure (e), vapor pressure deficit (VPD),
             specific humidity (q), and relative humidity (rH).

# Arguments
- Tair:      Air temperature (deg C)
- pressure:  Atmospheric pressure (kPa)
- e:         Vapor pressure (kPa)
- q:        Specific humidity (kg kg-1)
- VPD:       Vapor pressure deficit (kPa)
- rH:        Relative humidity (-)

All functions accept the optional arguemtns:
- `Esat_formula`: formula used in [`Esat_from_Tair`](@ref)
- `constants`: dictionary from [`BigleafConstants`](@ref) with entries
  eps and Pa2kPa

# Rreferences
Foken, T, 2008: Micrometeorology_ Springer, Berlin, Germany.
"""
function VPD_to_rH(VPD,Tair; Esat_formula=Sonntag1990(),
                      constants=BigleafConstants())
  esat = Esat_from_Tair(Tair; Esat_formula = Esat_formula,constants)
  rH   = 1 - VPD/esat
end,
function rH_to_VPD(rH,Tair; Esat_formula=Sonntag1990(),
                      constants=BigleafConstants())
  if !ismissing(rH) && rH > 1
    @warn("Expected relative humidity (rH) between 0 and 1, but was ", rH)
  end
  esat = Esat_from_Tair(Tair; Esat_formula = Esat_formula,constants)
  VPD  = esat - rH*esat
end,
function e_to_rH(e,Tair; Esat_formula=Sonntag1990(),
                    constants=BigleafConstants())
  esat = Esat_from_Tair(Tair; Esat_formula = Esat_formula,constants)
  if !ismissing(e) && (e > esat + sqrt(eps()))
    @warn("Provided vapour pressure that was higher than saturation_
             Returning rH=1 for those cases.")
  end
  rH  = min(1, e/esat)
end,
function VPD_to_e(VPD,Tair; Esat_formula=Sonntag1990(),
                     constants=BigleafConstants())
  esat = Esat_from_Tair(Tair; Esat_formula = Esat_formula,constants)
  e    = esat - VPD
end,
function e_to_VPD(e,Tair; Esat_formula=Sonntag1990(),
                     constants=BigleafConstants())
  esat = Esat_from_Tair(Tair; Esat_formula = Esat_formula,constants)
  VPD  = esat - e
end,
function e_to_q(e::FT,pressure; constants=BigleafConstants()) where FT
  q = (FT(constants.eps) * e / (pressure - (FT(1)-FT(constants.eps)) * e))
end,
function q_to_e(q::FT,pressure; constants=BigleafConstants()) where FT
  e = (q * pressure / ((FT(1)-FT(constants.eps)) * q + FT(constants.eps)))
end,
function q_to_VPD(q,Tair,pressure; Esat_formula=Sonntag1990(),
                     constants=BigleafConstants())
  esat = Esat_from_Tair(Tair; Esat_formula = Esat_formula,constants)
  e    = q_to_e(q,pressure; constants)
  VPD  = esat - e
end,
function VPD_to_q(VPD,Tair,pressure; Esat_formula=Sonntag1990(),
                     constants=BigleafConstants())
  esat = Esat_from_Tair(Tair; Esat_formula = Esat_formula,constants)
  e    = esat - VPD
  q    = e_to_q(e,pressure;constants)
end

"""
    Rg_to_PPFD(Rg,J_to_mol=4.6,frac_PAR=0.5)
    PPFD_to_Rg(PPFD,J_to_mol=4.6,frac_PAR=0.5)

Conversions between Global Radiation (W m-2) and Photosynthetic Photon Flux
Density (umol m-2 s-1)

# Arguments
- Rg:       Global radiation = incoming short-wave radiation at the surface (W m-2)
- PPFD:     Photosynthetic photon flux density (umol m-2 s-1)
- J_to_mol: Conversion factor from J m-2 s-1 (= W m-2) to umol (quanta) m-2 s-1
- frac_PAR: Fraction of incoming solar irradiance that is photosynthetical-
  active radiation (PAR); defaults to 0.5

# Details
The conversion is given by:

 ``PPFD = Rg * frac_PAR * J_to_mol``

by default, the combined conversion factor (`frac_PAR * J_to_mol`) is 2.3

# Examples
```@example
# convert a measured incoming short-wave radiation of 500 Wm-2 to
# PPFD in umol m-2 s-1 and backwards
Rg_to_PPFD(500)
PPFD_to_Rg(1150)
```
"""
function Rg_to_PPFD(Rg,J_to_mol=4.6,frac_PAR=0.5)
  PPFD = Rg * frac_PAR * J_to_mol
end,
function PPFD_to_Rg(PPFD,J_to_mol=4.6,frac_PAR=0.5)
  Rg = PPFD / frac_PAR / J_to_mol
end

"""
    kg_to_mol(mass, molarMass=BigleafConstants()[:H2Omol])

Conversion between Mass (kg) and Molar Units (mol).
"""
function kg_to_mol(mass, molarMass=BigleafConstants()[:H2Omol])
  moles = mass / molarMass
end

"""
    umolCO2_to_gC(CO2_flux; constants=BigleafConstants())
    gC_to_umolCO2(C_flux; constants=BigleafConstants())

Convert CO2 quantities from (umol CO2 m-2 s-1) to (g C m-2 d-1) and vice versa.

# Arguments
- CO2_flux  CO2 flux (umol CO2 m-2 s-1)
- C_flux    Carbon (C) flux (gC m-2 d-1)
- `constants`: dictionary from [`BigleafConstants`](@ref) with entries:
  Cmol, umol2mol, mol2umol, kg2g, g2kg, says2seconds

# Examples
```@example
umolCO2_to_gC(20)  # gC m-2 d-1
```
"""
function umolCO2_to_gC(CO2_flux::FT; constants=BigleafConstants()) where FT
  C_flux = CO2_flux * FT(constants.umol2mol) * FT(constants.Cmol) *
  FT(constants.kg2g) * FT(constants.days2seconds)
end,
function gC_to_umolCO2(C_flux::FT; constants=BigleafConstants()) where FT
  CO2_flux = (C_flux * FT(constants.g2kg)  / FT(constants.days2seconds)) /
  FT(constants.Cmol)  * FT(constants.mol2umol)
end


# Based on http://scikit-learn.org/stable/modules/gaussian_process.html
using PyCall
@pyimport numpy as np
@pyimport sklearn.gaussian_process as gausspr

f(x) = x .* sin(x)

xtrain = np.atleast_2d([0.9, 1.0, 3.0, 5.0, 7.0, 9.0, 9.1])'
ytrain = f(xtrain)

gp = gausspr.GaussianProcess(theta0=1e-2, thetaL=1e-4, thetaU=1e-1)
gp["fit"](xtrain, ytrain)

xtest = np.atleast_2d(np.linspace(0, 10, 1000))'
ytest = f(xtest)
y_pred, sigma2_pred = gp["predict"](xtest, eval_MSE=true)

mse(yhat, y) = mean(sum((yhat .- y).^2))
rmse(yhat, y) = sqrt(mse(yhat, y))
mape(yhat, y) = 100.0 * mean(abs( (y .- yhat)./y ))

rmse(y_pred[:], ytest[:])

# Around 1.0 it should be pretty close:
y_pred[95:105]
ytest[95:105]
mape(y_pred[95:105], ytest[95:105])


export MvNormalMeanCovariance

import Distributions: logdetcov, distrname, sqmahal, sqmahal!, AbstractMvNormal
import LinearAlgebra: diag, Diagonal, dot
import Base: ndims, precision, length, size, prod

struct MvNormalMeanCovariance{ T <: Real, M <: AbstractVector{T}, P <: AbstractMatrix{T} } <: AbstractMvNormal
    μ :: M
    Σ :: P
end

function MvNormalMeanCovariance(μ::AbstractVector{ <: Real }, Σ::AbstractMatrix{ <: Real })
    T = promote_type(eltype(μ), eltype(Σ))
    return MvNormalMeanCovariance(convert(AbstractArray{T}, μ), convert(AbstractArray{T}, Σ))
end

function MvNormalMeanCovariance(μ::AbstractVector{ <: Integer}, Σ::AbstractMatrix{ <: Integer })
    return MvNormalMeanCovariance(float.(μ), float.(Σ))
end

function MvNormalMeanCovariance(μ::AbstractVector, σ::AbstractVector)
    return MvNormalMeanCovariance(μ, matrix_from_diagonal(promote_type(eltype(μ), eltype(σ)), σ))
end

function MvNormalMeanCovariance(μ::AbstractVector{T}) where T
    return MvNormalMeanCovariance(μ, convert(AbstractArray{T}, ones(length(μ))))
end

Distributions.distrname(::MvNormalMeanCovariance) = "MvNormalMeanCovariance"

weightedmean(dist::MvNormalMeanCovariance) = invcov(dist) * mean(dist)

function weightedmean_invcov(dist::MvNormalMeanCovariance)
    W = invcov(dist)
    xi = W * mean(dist)
    return (xi, W)
end

Distributions.mean(dist::MvNormalMeanCovariance)      = dist.μ
Distributions.var(dist::MvNormalMeanCovariance)       = diag(cov(dist))
Distributions.cov(dist::MvNormalMeanCovariance)       = dist.Σ
Distributions.invcov(dist::MvNormalMeanCovariance)    = cholinv(dist.Σ)
Distributions.std(dist::MvNormalMeanCovariance)       = cholsqrt(cov(dist))
Distributions.logdetcov(dist::MvNormalMeanCovariance) = logdet(cov(dist))

Distributions.sqmahal(dist::MvNormalMeanCovariance, x::AbstractVector) = sqmahal!(similar(x), dist, x)

function Distributions.sqmahal!(r, dist::MvNormalMeanCovariance, x::AbstractVector)
    μ = mean(dist)
    for i in 1:length(r)
        @inbounds r[i] = μ[i] - x[i]
    end
    return xT_A_x(r, invcov(dist))
end

Base.eltype(::MvNormalMeanCovariance{T})     where T = T
Base.precision(dist::MvNormalMeanCovariance)         = invcov(dist)
Base.length(dist::MvNormalMeanCovariance)            = length(mean(dist))
Base.ndims(dist::MvNormalMeanCovariance)             = length(dist)
Base.size(dist::MvNormalMeanCovariance)              = (length(dist), )

Base.convert(::Type{ <: MvNormalMeanCovariance }, μ::AbstractVector, Σ::AbstractMatrix) = MvNormalMeanCovariance(μ, Σ)

function Base.convert(::Type{ <: MvNormalMeanCovariance{T} }, μ::AbstractVector, Σ::AbstractMatrix) where { T <: Real }
    return MvNormalMeanCovariance(convert(AbstractArray{T}, μ), convert(AbstractArray{T}, Σ))
end

vague(::Type{ <: MvNormalMeanCovariance }, dims::Int) = MvNormalMeanCovariance(zeros(dims), fill(huge, dims))

prod_analytical_rule(::Type{ <: MvNormalMeanCovariance }, ::Type{ <: MvNormalMeanCovariance }) = ProdAnalyticalRuleAvailable()

function Base.prod(::ProdPreserveType, left::MvNormalMeanCovariance, right::MvNormalMeanCovariance)
    invcovleft  = invcov(left)
    invcovright = invcov(right)
    Σ = cholinv(invcovleft + invcovright)
    μ = Σ * (invcovleft * mean(left) + invcovright * mean(right))
    return MvNormalMeanCovariance(μ, Σ)
end

function Base.prod(::ProdAnalytical, left::MvNormalMeanCovariance, right::MvNormalMeanCovariance)
    xi_left, W_left = weightedmean_precision(left)
    xi_right, W_right = weightedmean_precision(right)
    return MvNormalWeightedMeanPrecision(xi_left + xi_right, W_left + W_right)
end


include("../../__lib__/std/datastructs/src/datastructs_module.jl")
include("../../__lib__/std/util/src/util_module.jl")

module geom

using ..datastructs
using ..util
using LinearAlgebra
using Random
using PyPlot
PyPlot.pygui(true)

# =========================================================================== #
# Types
# =========================================================================== #
include("point.jl")
include("line.jl")
include("circle.jl")
include("quad.jl")
include("triangle.jl")
include("polygon.jl")

# =========================================================================== #
# Operations on types
# =========================================================================== #

# =========================================================================== #
# Mixed operations between types
# =========================================================================== #

# =========================================================================== #
# General, miscellenious types, operations
# =========================================================================== #
include("exterma.jl")
include("rotation_matrix.jl")
include("general.jl")

# =========================================================================== #
# Mesh generation
# =========================================================================== #
include("distance.jl")
include("convex_hull.jl")
include("delaunay.jl")
include("mesh_module.jl")

end # module


##
cd(@__DIR__())
using DelimitedFiles

##
path = "input.txt"
vals = readdlm(path, ' ', Int)

##
function find_triple(vals, res = 2020)
    for x in vals, y in vals, z in vals
        x + y + z == res && return x*y*z
    end
end

##
find_triple(vals)


aircraft = data_epema_rotor_sweep[1]
parameters = data_epema_rotor_sweep[2]

struct TestParams
    us
    vs
    wake_function
    wake_shape_functions
    axial_interpolations
    swirl_interpolations
    axial_multipliers
    swirl_multipliers
end

nsteps = length(parameters.us)
stepi = 5

testparams = TestParams(deepcopy(parameters.us), deepcopy(parameters.vs), Vector{Any}(nothing,nsteps), fill((Rtip, x) -> Rtip, length(aircraft.rotor_system.index)), fill((rs, us, r, Rtip) -> FM.linear(rs, us, r), length(aircraft.rotor_system.index)), fill((rs, vs, r, Rtip) -> FM.linear(rs, vs, r), length(aircraft.rotor_system.index)), fill((distance2plane, Rtip) -> 2, length(aircraft.rotor_system.index)), fill((distance2plane, Rtip) -> 1, length(aircraft.rotor_system.index)))

for i=1:nsteps
    AS.solve_rotor_wake(aircraft, testparams, nothing, nothing, nothing, i, nothing)
end

wakefun_aoa4 = testparams.wake_function[2]
r = aircraft.rotor_system.rotors[1].Rtip
# thetas = range(0, 2*pi, length=12)
# rs = range(r/10, 2*r, length=4)
# x = [r * cos(theta) for theta in thetas, r in rs]
# y = [r * cos(theta) for theta in thetas, r in rs]
# Vs = wakefun_aoa4.()
x = range(-r * 2.0, stop = r * 8.5, length=3)
y = range(-r * 2.5, stop = r * 2.5, length=35)
z = range(-r * 2.5, stop = r * 2.5, length=35)
lx = length(x)
ly = length(y)
lz = length(z)
xs = [x[i] for i in 1:lx, j in 1:ly, k in 1:lz]
ys = [y[j] for i in 1:lx, j in 1:ly, k in 1:lz]
zs = [z[k] for i in 1:lx, j in 1:ly, k in 1:lz]
Vs = wakefun_aoa4.([[x[i], y[j], z[k]] for i in 1:lx, j in 1:ly, k in 1:lz])
# Vs ./ AS.LA.norm.(Vs)
us = [Vs[i,j,k][1] for i in 1:lx, j in 1:ly, k in 1:lz]
vs = [Vs[i,j,k][2] for i in 1:lx, j in 1:ly, k in 1:lz]
ws = [Vs[i,j,k][3] for i in 1:lx, j in 1:ly, k in 1:lz]

fig = plt.figure("test_wakefunction")
fig.clear()
ax = fig.add_subplot(111, projection="3d")
ax.quiver3D(xs / r, ys / r, zs / r, us, vs, ws, length=0.01)#, normalize=true)

ax.set_xlabel(LS.L"x/r")
ax.set_ylabel(LS.L"y/r")
ax.set_zlabel(LS.L"z/r")

fig2d = plt.figure("test_wakefunction_2d")
fig2d.clear()
ax2 = fig2d.add_subplot(111)
ax2.quiver(ys[1,:,:] / r, zs[1,:,:] / r, (vs[1,:,:] .+ 1e-10), (ws[1,:,:] .+ 1e-10))#, length=0.01, normalize=true)
thetas = range(0,2*pi,length=100)
rotor_y = cos.(thetas)
rotor_z = sin.(thetas)
ax2.plot(rotor_y, rotor_z, color="blue")
ax2.set_xlabel(LS.L"x/r")
ax2.set_ylabel(LS.L"y/r")

rs = range(r/200, 1.5*r, length=100)
xs_3 = vcat(-reverse(rs), rs)
Vs_3 = wakefun_aoa4.([[0.0, r, 0.0] for r in xs_3])
us = [v[1] for v in Vs_3]
vs = [v[2] for v in Vs_3]
ws = [v[3] for v in Vs_3]

fig3 = plt.figure("test_wakefunction_us_vs")
fig3.clear()
ax_u = fig3.add_subplot(311)
ax_v = fig3.add_subplot(312)
ax_w = fig3.add_subplot(313)
ax_u.plot(xs_3 ./ r, us)
ax_u.set_ylabel(LS.L"u[m/s]")
ax_v.plot(xs_3 ./ r, vs)
ax_v.set_ylabel(LS.L"v[m/s]")
ax_w.plot(xs_3 ./ r, ws)
ax_w.set_ylabel(LS.L"w[m/s]")
ax_w.set_xlabel(LS.L"y/R, x=z=0")


using Test, FrameFun.Platforms, BasisFunctions

@testset "Generic Platforms" begin
    P = platform(Fourier(10))
    @test dictionary(P,param(Fourier(10))) == Fourier(10)
    @test P[10] == Fourier(10)
    @test operator(dualdictionary(P,10,measure(P))) ≈ operator(dual(Fourier(10)))
    @test DictionaryStyle(P)==BasisStyle()
    @test SamplingStyle(P) == InterpolationStyle()
    @test SolverStyle(P,SamplingStyle(P)) == InverseStyle()
    @test param_first(P) == 10
    @test param_double(P,10) == 20
    @test param_increment(P,10) == 11
    @test param_inbetween(P,10,20) == 15
    @test correctparamformat(P,10)
    @test !correctparamformat(P,(10,))

    @test elements(platform(Fourier(10)^2)) == (P,P)
    @test element(platform(Fourier(10)^2),1) == P

    P = platform(Fourier(10)^2)
    @test correctparamformat(P,(10,10))
    @test !correctparamformat(P,(10,))
    @test !correctparamformat(P,10)
    @test param_first(P) == (10,10)
    @test param_double(P,(10,20)) == (20,40)
    @test param_increment(P,(10,11)) == (11,12)
    @test param_inbetween(P,(10,10),(20,10)) == (15,10)
    @test ProductPlatform(platform(Fourier(10)),2) == P
    @test element(P,1) == platform(Fourier(10))
    length(elements(P))==2
    @test dictionary(P,10) == Fourier(10)^2
    @test dictionary(P,param(Fourier(10)^2)) == Fourier(10)^2
    @test P[10] == Fourier(10)^2
    @test dictionary(P,(11,12)) == Fourier(11)⊗Fourier(12)
    @test dualdictionary(P,(10,10), measure(Fourier(10)^2)) isa TensorProductDict
    isbasis(Fourier(10)^2)
    @test DictionaryStyle(P)==BasisStyle()
    @test SamplingStyle(P) == ProductSamplingStyle(InterpolationStyle(),InterpolationStyle())
    @test SolverStyle(P,SamplingStyle(P)) == ProductSolverStyle(InverseStyle(),InverseStyle())
end

@testset "Platform Styles" begin
    @test SamplingStyle(platform(Fourier(10))) == InterpolationStyle()
    @test SolverStyle(platform(Fourier(10)), InterpolationStyle()) == InverseStyle()
    @test SolverStyle(platform(Fourier(10)), OversamplingStyle()) == DirectStyle()
    @test SolverStyle(platform(Fourier(10)), GridStyle()) == DirectStyle()
    @test SolverStyle(platform(Fourier(10)), GenericSamplingStyle()) == DirectStyle()

    @test SolverStyle(FrameStyle(), platform(Fourier(10)), InterpolationStyle()) == AZStyle()
    @test SolverStyle(FrameStyle(), platform(Fourier(10)), OversamplingStyle()) == AZStyle()
    @test SolverStyle(FrameStyle(), platform(Fourier(10)), GridStyle()) == AZStyle()
    @test SolverStyle(FrameStyle(), platform(Fourier(10)), GenericSamplingStyle()) == AZStyle()

    @test DictionaryStyle(platform(Fourier(10))) == BasisStyle()
    @test DictionaryStyle(platform(MultiDict((Fourier(10),Fourier(10))))) == UnknownDictionaryStyle()
    @test elements(SolverStyle(platform(Fourier(10)^2), SamplingStyle(platform(Fourier(10)^2)))) == (InverseStyle(), InverseStyle())
    @test ProblemStyle(platform(Fourier(10)))  == DictionaryOperatorStyle()

end


#nprocs() == 1 && addprocs()

using GZip
using MNIST
using Knet
@everywhere using Compat

@everywhere logis(x) = 1 / (1 + exp(-x))
@everywhere dlogis(x) = logis(x) * (1 - logis(x))

#@everywhere logis(x) = max(0,x)
#@everywhere dlogis(x) = x>0?1:0
@everywhere rmse(x) = sqrt(mean(x.^2))

@everywhere function initParams(x,y,n_hidden,max_iter)
    global w1 = randn(n_hidden,size(x,1))
    global w2 = randn(n_hidden,n_hidden)
    global w3 = randn(size(y,1),n_hidden)
    B1 = rand(n_hidden,size(y,1))
    B2 = rand(n_hidden,size(y,1))
    errors = zeros(max_iter)
    outputs = zeros(size(y,2),max_iter)
    return (w1,w2,w3,errors,outputs,B1,B2)
end

@everywhere function ffward(x,y,w1,w2,w3)
    a1 = w1 * x
    h1 = logis.(a1)
    a2 = w2 * h1
    h2 = logis.(a2)
    ay = w3 * h2
    yhat = logis.(ay)
    err = yhat - y
    return(a1,h1,a2,h2,ay,yhat,err)
end

@everywhere function updateW!(w1,w2,w3,x,h1,h2,δa1,δa2,err;lr=0.1)
    w1 .-= δa1 * x' .* lr
    w2 .-= δa2 * h1' .* lr
    w3 .-= err * h2' .* lr
end

@everywhere function bp(dta,n_hidden,max_iter)
    (x,y)=dta[1]
    (w1,w2,w3,errors,outputs,_,_)=initParams(x,y,n_hidden,max_iter)
    for i = 1:max_iter
        toterr=0
        for (x,y) in dta
            (a1,h1,a2,h2,ay,yhat,err) = ffward(x,y,w1,w2,w3)
            δa2 = (w3' * err) .* dlogis.(a2)
            δa1 = (w2' * δa2) .* dlogis.(a1)
            updateW!(w1,w2,w3,x,h1,h2,δa1,δa2,err)
            toterr += rmse(err)
        end
        println(toterr)
        errors[i] = toterr/length(dta)
    end
    return errors
end

@everywhere function dfa(dta,n_hidden,max_iter)
    (x,y)=dta[1]
    (w1,w2,w3,errors,outputs,B1,B2)=initParams(x,y,n_hidden,max_iter)
    for i = 1:max_iter
        toterr=0
        for (x,y) in dta
            (a1,h1,a2,h2,ay,yhat,err) = ffward(x,y,w1,w2,w3)
            δa2 = (B2*err) .* dlogis.(a2)
            δa1 = (B1*err) .* dlogis.(a1)
            updateW!(w1,w2,w3,x,h1,h2,δa1,δa2,err)
            toterr += rmse(err)
        end
        println(toterr)
        errors[i] = toterr/length(dta)
    end
    return errors
end

@everywhere function fa(x,y,n_hidden,max_iter)
    (w1,w2,w3,errors,outputs,B1,B2)=initParams(x,y,n_hidden,max_iter)
    for i = 1:max_iter
        (a1,h1,a2,h2,ay,yhat,err) = ffward(x,y,w1,w2,w3)
        δa2 = (B2.*err) .* dlogis(a2)
        δa1 = (B1.*δa2) .* dlogis(a1)
        updateW!(w1,w2,w3,x,h1,h2,δa1,δa2,err)
        errors[i] = sum(abs(err))
        outputs[:,i] = yhat
    end
    return (outputs,errors)
end

@everywhere function ifa(x,y,n_hidden,max_iter)
    (w1,w2,w3,errors,outputs,B1,_)=initParams(x,y,n_hidden,max_iter)
    for i = 1:max_iter
        (a1,h1,a2,h2,ay,yhat,err) = ffward(x,y,w1,w2,w3)
        δa1 = (B1.*err) .* dlogis(a1)
        δa2 = (w2 * δa1) .* dlogis(a2)
        updateW!(w1,w2,w3,x,h1,h2,δa1,δa2,err)
        errors[i] = sum(abs(err))
        outputs[:,i] = yhat
    end
    return (outputs,errors)
end

@everywhere function loaddata()
    global xtrn,ytrn,xtst,ytst
    info("Loading MNIST...")
    xtrn = gzload("train-images-idx3-ubyte.gz")[17:end]
    xtst = gzload("t10k-images-idx3-ubyte.gz")[17:end]
    ytrn = gzload("train-labels-idx1-ubyte.gz")[9:end]
    ytst = gzload("t10k-labels-idx1-ubyte.gz")[9:end]
end

@everywhere function gzload(file; path=Knet.dir("data",file), url="http://yann.lecun.com/exdb/mnist/$file")
    isfile(path) || download(url, path)
    f = gzopen(path)
    a = @compat read(f)
    close(f)
    return(a)
end

@everywhere function minibatch(x, y, batchsize; atype=Array{Float32}, xrows=784, yrows=10, xscale=255)
    xbatch(a)=convert(atype, reshape(a./xscale, xrows, div(length(a),xrows)))
    ybatch(a)=(a[a.==0]=10; convert(atype, sparse(convert(Vector{Int},a),1:length(a),one(eltype(a)),yrows,length(a))))
    xcols = div(length(x),xrows)
    xcols == length(y) || throw(DimensionMismatch())
    data = Any[]
    for i=1:batchsize:xcols-batchsize+1
        j=i+batchsize-1
        push!(data, (xbatch(x[1+(i-1)*xrows:j*xrows]), ybatch(y[i:j])))
    end
    return data
end

loaddata()
global dtrn = minibatch(xtrn, ytrn, 100; atype=Array{Float32})
global dtst = minibatch(xtst, ytst, 100; atype=Array{Float32})

#x = [1 1 ; 0 1 ; 1 0 ; 0 0]'
#y = [0 1 1 0]

n_hidden = 50; # Number of hidden units
max_iter = 10; #Number of learning steps
trials = 1;

#dfa_err=@parallel (hcat) for _=1:trials
#    dfa_errors = dfa(dtrn,n_hidden,max_iter)
#    dfa_errors
#end

bp_err = @parallel (hcat) for _=1:trials
    bp_errors = bp(dtrn,n_hidden,max_iter)
    bp_errors
end

#(bp_yhat,_) = bp(x,y,n_hidden,max_iter)
#
#ifa_err=@parallel (hcat) for _=1:trials
#    (_,ifa_errors) = ifa(x,y,n_hidden,max_iter)
#    ifa_errors
#end
#(ifa_yhat,_) = ifa(x,y,n_hidden,max_iter);
#fa_err=@parallel (hcat) for _=1:trials
#    (_,fa_errors) = fa(x,y,n_hidden,max_iter)
#    fa_errors
#end
#(fa_yhat,_) = fa(x,y,n_hidden,max_iter);
#using Plots
#l = @layout [
#    a b
#    c d
#    e f
#]

#unicodeplots()
##init_notebook(true)
#p1=plot(dfa_yhat',label=["1,1" "0,1" "1,0" "0,0"],title="DFA")
#p2=plot(ifa_yhat',label=["1,1" "0,1" "1,0" "0,0"],title="IFA")
#p3=plot(fa_yhat',label=["1,1" "0,1" "1,0" "0,0"],title="FA")
#p4=plot(bp_yhat',label=["1,1" "0,1" "1,0" "0,0"],title="BP")
#p5=plot([dfa_err[:,1] ifa_err[:,1] fa_err[:,1] bp_err[:,1]],label=["DFA" "IFA" "FA" "BP"],title="First Trial Error")
#p6=plot([mean(dfa_err,2) mean(ifa_err,2) mean(fa_err,2) mean(bp_err,2)],label=["DFA" "IFA" "FA" "BP"],title="Mean Trials Error")
#plot(p1,p2,p3,p4,p5,p6,layout=l)
##plot(p1,p2,p3,p4,p5,p6,size=(900,900),layout=l)


#=

MuSSE equations

Ignacio Quintero Mächler

t(-_-t)

September 16 2017

=#

"""
    make_musse(t,u,p,du)

MuSSE equation function for integration.
"""
function make_musse(k::Int64)

  Qf! = makeQ(k)

  # preallocate matrices & vectors
  Q ::Array{Float64,2} = zeros(k, k)
  Qs::Array{Float64,1} = zeros(k)
  Qe::Array{Float64,1} = zeros(k)

  # estimate indices previously
  qsi::UnitRange{Int64} = 1:k
  qei::UnitRange{Int64} = (k+1):2k

  function f(du, u, p, t)

    Qf!(Q, p)
    BLAS.gemv!('N', 1.0, Q, u[qsi], 0.0, Qs)
    BLAS.gemv!('N', 1.0, Q, u[qei], 0.0, Qe)

    for i in Base.OneTo(k)
      du[i]   =        - (p[i] + p[k+i])*u[i]   + 2.0*p[i]*u[i+k]*u[i]   + Qs[i]
      du[i+k] = p[k+i] - (p[i] + p[k+i])*u[i+k] +     p[i]*u[i+k]*u[i+k] + Qe[i]
    end

  end

  return f
end

"""
    makeQ(k::Int64)

Transform instantaneous transition rates into the instantaneous
transition matrix `Q`.
"""
function makeQ(k::Int64)

  pidx::UnitRange{Int64}  = (2k+1):(k + k*k)
  Didx::StepRange{Int64}  = 1:(k+1):(k*k)
  Qidx::Array{Int64,1}    = setdiff(1:k*k, Didx)
  pitr::Base.OneTo{Int64} = Base.OneTo(lastindex(pidx))
  kitr::Base.OneTo{Int64} = Base.OneTo(k)

  function f(Q::Array{Float64,2}, p::Array{Float64,1})

    @inbounds begin
      # assign parameters
      @simd for i in pitr
        Q[Qidx[i]] = p[pidx[i]]
      end

      #assign 0s
      for i in Didx
        Q[i] = 0.0
      end

      #row sums
      for j = kitr, i = kitr
        if i == j
          continue
        end
        Q[Didx[i]] -= Q[i,j]
      end
    end

  return nothing
  end
end


using Distances, LinearAlgebra
using Test
using KernelFunctions
using PDMats

dims = [10,5]

A = rand(dims...)
B = rand(dims...)
C = rand(8,9)
K = [zeros(dims[1],dims[1]),zeros(dims[2],dims[2])]
Kdiag = [zeros(dims[1]),zeros(dims[2])]
k = SqExponentialKernel()
@testset "Kernel Matrix Operations" begin
    @testset "Inplace Kernel Matrix" begin
        for obsdim in [1,2]
            @test kernelmatrix!(K[obsdim],k,A,B,obsdim=obsdim) == kernelmatrix(k,A,B,obsdim=obsdim)
            @test kernelmatrix!(K[obsdim],k,A,obsdim=obsdim) == kernelmatrix(k,A,obsdim=obsdim)
            @test kerneldiagmatrix!(Kdiag[obsdim],k,A,obsdim=obsdim) == kerneldiagmatrix(k,A,obsdim=obsdim)
            @test_throws DimensionMismatch kernelmatrix!(K[obsdim],k,A,C,obsdim=obsdim)
            @test_throws DimensionMismatch kernelmatrix!(K[obsdim],k,C,obsdim=obsdim)
            @test_throws DimensionMismatch kerneldiagmatrix!(Kdiag[obsdim],k,C,obsdim=obsdim)
        end
    end
    @testset "Kernel matrix" begin
        for obsdim in [1,2]
            @test kernelmatrix(k,A,B,obsdim=obsdim) == kappa.(k,pairwise(KernelFunctions.metric(k),A,B,dims=obsdim))
            @test kernelmatrix(k,A,obsdim=obsdim) == kappa.(k,pairwise(KernelFunctions.metric(k),A,dims=obsdim))
            @test kerneldiagmatrix(k,A,obsdim=obsdim) == diag(kernelmatrix(k,A,obsdim=obsdim))
            @test k(A,B,obsdim=obsdim) == kernelmatrix(k,A,B,obsdim=obsdim)
            @test k(A,obsdim=obsdim) == kernelmatrix(k,A,obsdim=obsdim)
            @test KernelFunctions._kernel(k,1.0,2.0) == KernelFunctions._kernel(k,[1.0],[2.0])
            @test_throws DimensionMismatch kernelmatrix(k,A,C,obsdim=obsdim)
        end
    end
    @testset "KernelSum" begin
        k1 = SqExponentialKernel()
        k2 = LinearKernel()
        ks = k1 + k2
        w1 = 0.4; w2 = 1.2;
        ks2 = KernelSum([k1,k2],weights=[w1,w2])
        @test all(kernelmatrix(ks,A) .== kernelmatrix(k1,A) + kernelmatrix(k2,A))
        @test all(kernelmatrix(ks+k1,A) .≈ 2*kernelmatrix(k1,A) + kernelmatrix(k2,A))
        @test all(kernelmatrix(k1+ks,A) .≈ 2*kernelmatrix(k1,A) + kernelmatrix(k2,A))
        @test all(kernelmatrix(ks,A,B) .== kernelmatrix(k1,A,B) + kernelmatrix(k2,A,B))
        @test all(kerneldiagmatrix(ks,A) .== kerneldiagmatrix(k1,A) + kerneldiagmatrix(k2,A))
        @test all(kernelmatrix(ks2,A) .== w1*kernelmatrix(k1,A) + w2*kernelmatrix(k2,A))
    end
    @testset "KernelProduct" begin
        k1 = SqExponentialKernel()
        k2 = LinearKernel()
        k3 = RationalQuadraticKernel()
        kp = k1 * k2
        kp2 = k1 * k3
        @test all(KernelFunctions.metric(kp).==[KernelFunctions.metric(k1),KernelFunctions.metric(k2)])
        @test all(kernelmatrix(kp,A) .≈ kernelmatrix(k1,A) .* kernelmatrix(k2,A))
        @test all(kernelmatrix(kp*k1,A) .≈ kernelmatrix(k1,A).^2 .* kernelmatrix(k2,A))
        @test all(kernelmatrix(k1*kp,A) .≈ kernelmatrix(k1,A).^2 .* kernelmatrix(k2,A))
        @test all(kernelmatrix(kp,A) .≈ kernelmatrix(k1,A) .* kernelmatrix(k2,A))
        @test all(kernelmatrix(kp,A,B) .≈ kernelmatrix(k1,A,B) .* kernelmatrix(k2,A,B))
        @test all(kernelmatrix(kp,A) .≈ kernelmatrix(k1,A) .* kernelmatrix(k2,A))
        @test all(kerneldiagmatrix(kp,A) .== kerneldiagmatrix(k1,A) .* kerneldiagmatrix(k2,A))
    end
    @testset "PDMat" begin
        for obsdim in [1,2]
            @test all(Matrix(kernelpdmat(k,A,obsdim=obsdim)) .≈ Matrix(PDMat(kernelmatrix(k,A,obsdim=obsdim))))
            # @test_throws ErrorException kernelpdmat(k,ones(100,100),obsdim=obsdim)
        end
    end
end


using WAV

"""
### waverender(Timbrer;frequency, framerate=44100, amplitude=0.5)
- Timbrer as the timbre function, with T as 2pi, amplitude as 1.

return the standard wave synthesis method
"""
function waverender(Timbrer;frequency=440.0, framerate=44100, amplitude=0.5)
    amplitude > 1 && (amplitude = 1.0)
    amplitude < 0 && (amplitude = 0.0)
    return (i) -> begin
        T = 1.0 / frequency
        t = 1.0 / framerate * i
        amplitude * Timbrer(2.0 * pi * t / T)
    end
end;

"""
### computesamples(noteArray::AbstractArray;framerate=44100,shader=nothing::Any,shaderportion=0.1)
- noteArray as Array{[wavefunction,ticklen]}
- shader should be a function, take 1 float parameter, range from 0 to 1. return a scale from 1 to 0.

return the rendered data, as Array{Float64, 1}
"""
function computesamples(noteArray::AbstractArray;framerate=44100.0,shader=nothing::Any,shaderportion=0.1)
    finalRender = Array{Float64}(0)
    shader == nothing && (shader = (x) -> 1 - x)
    for (wavefunction,ticklen) in noteArray
        shaderStart = ticklen * ( 1.0 - shaderportion)
        finalRender = vcat(finalRender, Float32[wavefunction(i) for i = 1:shaderStart])
        finalRender = vcat(finalRender, Float32[wavefunction(i) * shader((i-shaderStart)/(ticklen-shaderStart)) for i = shaderStart:ticklen])
    end
    return finalRender
end;

toneDict = Dict(
    "C" => 1,
    "C#" => 2,
    "D" => 3,
    "D#" => 4,
    "E" => 5,
    "F" => 6,
    "F#" => 7,
    "G" => 8,
    "G#" => 9,
    "A" => 10,
    "A#" => 11,
    "B" => 12
);

function composer(sheet::String; f0 = 440.0, unit = 1.0, gap=0.05, timbrer=sin, framerate=44100.0, base_amplitude=0.5, shader=nothing,shaderportion=0.1)
    noteArray = Array{Array}(0)
    for note in split(sheet)
        if note == "-"
            noteArray[end-1][end] += unit*framerate
            continue
        end
        m = match(r"(?<tone>\w#?)(?<octave>\d)",note)
        tone = m[:tone]
        octave = Int(parse(m[:octave]))
        frequency = f0 * 2 ^(((toneDict[tone] + (octave-4)*12) - toneDict["A"])/12)

        noteArray = vcat(noteArray, [[waverender(timbrer,frequency=frequency),(1-gap)*framerate*unit],[waverender(timbrer,amplitude=0.0),gap*framerate*unit]])
    end
    render = computesamples(noteArray,framerate=framerate,shader=shader,shaderportion=shaderportion);
    return render
end;

function composer(sheet::String, outfilename::String; f0 = 440.0, unit=1.0, gap=0.05, timbrer=sin, framerate=44100.0, base_amplitude=0.5, shader=nothing,shaderportion=0.1)
    samples = composer(sheet, f0=f0, unit=unit,gap=gap, timbrer=timbrer, framerate=framerate, base_amplitude=base_amplitude, shader=shader,shaderportion=shaderportion)
    wavwrite(samples, outfilename, Fs=framerate);
    return samples
end;


# Using Base modules.
using Random
# Load a plotting library.
using Plots
# Load the distributions library.
using Distributions

using StatsPlots

mutable struct CAVI <: InferenceAlgorithm
    # n_iters   ::  Int       # maximum number of iterations
    X # Data
    K::Int       # number of iterations
    N::Int
    sigma2
    # tol :: T    # tolerence values for early convergence criteria
    m
    phi
    s2
    # proposals ::  Dict{Symbol,Any}  # Proposals for paramters
    # space     ::  Set{T}    # sampling space, emtpy means all
end

function CAVI(X, num_components::Int, n_samples::Int, sigma2=1)
    K=num_components
    N=n_samples
    # global sigma2 = 1
    rng = RandomDevice()
    phi = rand(Dirichlet(K*N, (rand(rng, Float32,(1))*rand(rng, 1:9))[1]),(3));
    m = rand(rng,floor(minimum(X)):floor(maximum(X)), (1,K))
    m += maximum(X)*rand(rng, Float32, (1,K))
    s2 = rand(rng, Float32, (1, K))
    print("\nInit mean\n")
    println(m)
    print("\nInit s2\n")
    println(s2)
    CAVI(X, num_components, n_samples, sigma2, m, phi, s2)

end

function get_elbo(alg::CAVI)
    t1 = log.(alg.s2) - alg.m./(alg.sigma2)
    t1 = sum(t1)
    t2 = -0.5.*((alg.X).^2 * ones(1, 3) + ones(3000, 1)*((alg.s2)+(alg.m).^2 ))
    t2 .+= (alg.X)*(alg.m)
    t2 .-= log.(alg.phi)
    t2 .+= alg.phi
    t2 = sum(t2)
    return t1 + t2
end

function _update_phi(alg::CAVI)
    t1 = alg.X*alg.m
    t2 = -(0.5*alg.m.^2 + 0.5*alg.s2)
    exponent = t1 .+ t2
    phi = exp.(exponent)
    alg.phi = phi ./ sum(phi,dims=2)
end
function _update_mu(alg::CAVI)
    alg.m = sum((alg.phi.*alg.X),dims=1) .* (1 ./alg.sigma2 .+ sum(alg.phi,dims=1)).^(-1)
    @assert size(alg.m)[2] == alg.K
    #print(self.m)
    alg.s2 = (1 ./alg.sigma2 .+ sum(alg.phi,dims=1)).^(-1)
    @assert size(alg.s2)[2] == alg.K
end

function _cavi_step(alg::CAVI)
    _update_phi(alg)
    _update_mu(alg)
end

function fit(alg::CAVI; max_iter=100, tol=1e-20, print_interval=50)
    # _init(X=X, num_components=num_components, n_samples=n_samples, sigma2=sigma2)
    elbo_values = [get_elbo(alg)]
    m_history = [alg.m]
    s2_history = [alg.s2]
    for iter_ in 1:(max_iter+1)
        _cavi_step(alg)
        append!(m_history, [alg.m])
        append!(s2_history, [alg.s2])
        append!(elbo_values, [get_elbo(alg)])

        if iter_ % print_interval == 0
            println(iter_, " Means: ", m_history[iter_], " ELBO: ", elbo_values[end])
        end
        if abs(elbo_values[end-1] - elbo_values[end]) < tol
            i=iter_
            elbo_val = elbo_values[end]
            println("ELBO converged with $elbo_val at iteration $i")
            break
        end
        if iter_ == max_iter
            elbo_val = elbo_values[end]
            print("ELBO ended with ll $elbo_val")
        end
    end
    elbo_values, m_history ,s2_history
end


"""
    A = JopFft(T, n...[; dims=()])

where `A` is a Fourier transform, the domain of A is (T,n).  If `dims`
is not set, then `A` is an N-dimensional Fourier transform and where `N=ndims(dom)`.
Otherwise, `A` is a Fourier transform over the dimensions specifed by `dims`.

For a complex-to-complex transform: `range(A)::JetSpace`.  For a real-to-complex
transform: `range(A)::JetSpaceSymmetric`.  Normally, JopFft is responsible for constructing
range(A).  In the event that you need to manually construct this space, we provide
a convenience method:

    R = symspace(JopLn{JopFft_df!}, T, symdim, n...)

where `T` is the precision (generally `Float32` or `Float64`), `n...` is
the dimensions of the domain of `A`, and `symdim` is the first dimension
begin Fourier transformed.

# Examples

## 1D, real to complex
```julia
A = JopFft(Float64,128)
m = rand(domain(A))
d = A*m
```
## 1D, complex to complex
```julia
A = JopFft(ComplexF64,128)
m = rand(domain(A))
d = A*m
```

## 2D, real to complex
```julia
A = JopFft(Float64,128,256)
m = rand(domain(A))
d = A*m
```

## 3D, real to complex, transforming over only the first dimension
```julia
A = JopFft(Float64,128,256; dims=(1,))
m = rand(domain(A))
d = A*m
```
"""
function JopFft(::Type{T},n::Vararg{Int,N};dims=()) where {T<:Real,N}
    dims = isempty(dims) ? ntuple(i->i, N) : dims
    JopLn(dom = JetSpace(T,n), rng = symspace(typeof(JopFft_df!),T,dims[1],n...), df! = JopFft_df!, df′! = JopFft_df′!, s = (dims=dims,))
end
function JopFft(::Type{T},n::Vararg{Int,N};dims=()) where {T<:Complex,N}
    dims = isempty(dims) ? ntuple(i->i, N) : dims
    JopLn(dom = JetSpace(T,n), rng = JetSpace(T,n), df! = JopFft_df!, df′! = JopFft_df′!, s = (dims=dims,))
end
JopFft(R::JetSpace; dims=()) = JopFft(eltype(R), size(R)...; dims=dims)
export JopFft

function JopFft_df!(d::Jets.SymmetricArray{Complex{T}}, m::AbstractArray{T}; kwargs...) where {T}
    JopFft_df!(parent(d), m; kwargs...)
    d
end
function JopFft_df!(d::Array{Complex{T}}, m::AbstractArray{T}; dims, kwargs...) where {T<:Real}
    P = plan_rfft(m, dims)
    JopFft_mul!(d, P, m, dims)
    d
end

function JopFft_df!(d::AbstractArray{T}, m::AbstractArray{T}; dims, kwargs...) where {T<:Complex}
    P = plan_fft(m, dims)
    JopFft_mul!(d, P, m, dims)
    d
end

function JopFft_mul!(d::AbstractArray{T}, P::FFTW.FFTWPlan, m::AbstractArray, dims) where {T}
    mul!(d, P, m)
    sc = fftscale(T, size(m), dims)
    d .*= sc
end

JopFft_df′!(m::AbstractArray{T}, d::Jets.SymmetricArray{Complex{T}}; kwargs...) where {T} = JopFft_df′!(m, parent(d); kwargs...)
function JopFft_df′!(m::AbstractArray{T}, d::Array{Complex{T}}; dims, kwargs...) where {T<:Real}
    P = plan_brfft(d, size(m, dims[1]), dims)
    JopFft_mul_adjoint!(m, P, copy(parent(d)), dims) # in-place version of brfft modifies the input buffer
    m
end

function JopFft_df′!(m::AbstractArray{T}, d::AbstractArray; dims, kwargs...) where {T<:Complex}
    P = plan_bfft(d, dims)
    JopFft_mul_adjoint!(m, P, d, dims)
    m
end

function JopFft_mul_adjoint!(m::AbstractArray{T}, P::FFTW.FFTWPlan, d::AbstractArray, dims) where {T}
    mul!(m, P, d)
    sc = fftscale(T, size(m), dims)
    m .*= sc
end

fftscale(::Type{T}, n, dims) where {T} = T(1.0/sqrt(mapreduce(dim->n[dim], *, dims)))

function _JopFft_symspace_map(idim, I, symdim, n)
    if idim == symdim
        return 2 + n[idim] - I[idim]
    else
        return I[idim]
    end
end

function JopFft_symspace_map(I, symdim::Int, n::NTuple{N,Int}, _n::NTuple{N,Int}) where {N}
    if I[symdim] > _n[symdim]
        J = ntuple(idim->_JopFft_symspace_map(idim, I, symdim, n), N)
        return CartesianIndex(J)
    else
        return CartesianIndex(I)
    end
end

function Jets.symspace(JopFft, ::Type{T}, symdim::Int, n::Vararg{Int,N}) where {T,N}
    _n = ntuple(idim->idim == symdim ? div(n[idim], 2) + 1 : n[idim], N)
    Jets.JetSSpace(Complex{T}, n, _n, I->JopFft_symspace_map(I, symdim, n, _n))
end


module Synchrony
using DSP

export frequencies

# Get next fast FFT size for a given signal length
nextfastfft(n) = nextprod([2, 3, 5, 7], n)

# Compute frequencies based on length of padded FFT
function frequencies(nfft::Int, fs::Real=1.0)
    freq = (0:div(nfft, 2))*(fs/nfft)
    (freq, 1:length(freq))
end

# Compute frequencies based on length of padded FFT, with limits
function frequencies(nfft::Int, fs::Real, fmin::Real, fmax::Real=Inf)
    (freq,) = frequencies(nfft, fs)
    freqrange = searchsortedfirst(freq, fmin):(fmax == Inf ? length(freq) : searchsortedlast(freq, fmax))
    (freq[freqrange], freqrange)
end

# Compute frequencies based on data length, assuming padding to nextfastfft
frequencies(A::Union{AbstractVector{T}, AbstractMatrix{T}, AbstractArray{T,3}},
            fs::Real=1.0, fmin::Real=0.0, fmax::Real=Inf) where {T<:Real} =
    frequencies(nextfastfft(size(A, 1)), fs, fmin, fmax)

# Get the equivalent complex type for a given type
complextype(::Type{T}) where {T<:Complex} = T
complextype(::Type{T}) where {T} = Complex{T}

# Get preferred output type for a given input type
outputtype(::Type{T}) where {T<:AbstractFloat} = T
outputtype(::Type{T}) where {T<:Real} = Float64

include("transform_stats.jl")
include("point_field.jl")
include("wavelet.jl")
end


function local_richardson(uvT, 𝒟; unscale)
    H = Float32(abs(𝒟.uw.z[end] - 𝒟.uw.z[1]))
    g = 9.81f0
    α = 1.67f-4
    u_scaling = 𝒟.scalings["u"]
    v_scaling = 𝒟.scalings["v"]
    T_scaling = 𝒟.scalings["T"]
    σ_u = Float32(u_scaling.σ)
    σ_v = Float32(v_scaling.σ)
    σ_T = Float32(T_scaling.σ)

    Nz = Int(size(uvT, 1) / 3)
    D_cell = Float32.(Dᶜ(Nz, 1 / Nz))
    D_face = Float32.(Dᶠ(Nz, 1 / Nz))
    Ris = similar(uvT, Nz + 1, size(uvT,2))
    for i in 1:size(Ris, 2)
        u = @view uvT[1:Nz, i]
        v = @view uvT[Nz + 1:2Nz, i]
        T = @view uvT[2Nz + 1:3Nz, i]
        if unscale
            u .= u_scaling.(u)
            v .= v_scaling.(v)
            T .= T_scaling.(T)
        end
        Ri = @view Ris[:, i]
        ∂u∂z = D_face * u
        ∂v∂z = D_face * v
        ∂T∂z = D_face * T
        Ri .= (H * g * α * σ_T .* ∂T∂z) ./ ((σ_u .* ∂u∂z) .^2 + (σ_v .* ∂v∂z) .^2)
    end

    for i in 1:length(Ris)
        if isnan(Ris[i])
            Ris[i] = 0
        end

        if Ris[i] == Inf
            Ris[i] = 10f10
        end
    end
    return Ris
end


module VorticityStream

using LinearAlgebra
using Base.Iterators
using StaticArrays
using Statistics
using TimerOutputs

using ..AeroMDAO: AbstractPanel, AbstractPanel2D, Panel2D, WakePanel2D, Point2D, Uniform2D, collocation_point, p1, p2, transform_panel, transform_panel_points, affine_2D, panel_length, panel_angle, panel_tangent, panel_normal, distance, panel_points, stream, pressure_coefficient, rotation, inverse_rotation, midpair_map, panel_points, reverse_panel, panel_scalar, trailing_edge_panel

include("singularities.jl")
include("matrix.jl")

## Dynamics helpers
#===========================================================================#

lift_coefficient(cp, dist_colpoints, panel_angle) = - cp * dist_colpoints * cos(panel_angle)

lift_coefficient(wake_strength, speed) = 2. * wake_strength / speed

## Matrix assembly
#===========================================================================#

# export solve_problem

# function solve_problem(panels :: Vector{<: Panel2D}, u, sources :: Bool, wake_length)
#   φs        = solve_linear(panels, u, sources; bound = wake_length)
#   cps, cls  = aerodynamic_coefficients(panels, φs, u, sources)
#   cl_wake   = lift_coefficient(last(φs), norm(u))

#   cps, cls, cl_wake
# end

# function solve_problem(panels :: Vector{<: Panel2D}, u, num_wake :: Integer, wake_length)
#   wakes     = wake_panels(panels, wake_length, num_wake)
#   φs    = solve_linear(panels, u, wakes; bound = wake_length)
#   cps, cls  = aerodynamic_coefficients(panels, φs, u, true)
#   cl_wake   = lift_coefficient(last(φs), norm(u))

#   cps, cls, cl_wake
# end

# function ψp_influence(panel_i :: Panel2D, point)
#   xp, yp = transform_panel(panel_i, point)
#   vortex_stream_1(1., xp, yp, 0., panel_length(panel_i)) + vortex_stream_2(1., xp, yp, 0., panel_length(panel_i))
# end

# function ψm_influence(panel_i :: Panel2D, point)
#   xp, yp = transform_panel(panel_i, point)
#   vortex_stream_1(1., xp, yp, 0., panel_length(panel_i)) - vortex_stream_2(1., xp, yp, 0., panel_length(panel_i))
# end

# function te_influence(panel_i :: Panel2D, point)
#   xp, yp = transform_panel(panel_i, point)
#   source_stream(1., xp, yp, 0., panel_length(panel_i)) + vortex_stream_1(1., xp, yp, 0., panel_length(panel_i))
# end

# function source_influence(panel_i :: Panel2D, point)
#   xp, yp = transform_panel(panel_i, point)
#   source_stream(1., xp, yp, 0., panel_length(panel_i))
# end

# function source_influence_te(panel_i :: Panel2D, point)
#   xp, yp = transform_panel(panel_i, point)
#   source_stream_te(1., xp, yp, 0., panel_length(panel_i))
# end

end


@testset "Low Pass Tests" begin
    @testset "Agreement with explicit smoothing for random data" begin
        ts = cumsum(rand(1000))
        ys = randn(1000)
        dys = 0.1 .+ rand(1000)

        i = 100
        tau = pi

        wts = exp.(-abs.(ts.-ts[i])./tau)
        yout = sum(ys.*wts)/sum(wts)
        dyout = sqrt(sum(wts.*wts.*dys.*dys)/(sum(wts)*sum(wts)))

        ys_filt, dys_filt = LowPass.exp_lowpass(ts, ys, dys, tau)

        @test abs(ys_filt[i] - yout) < 1e-10
        @test abs(dys_filt[i] - dyout) < 1e-10
    end
end


"""
    BSplineBasis{T<:AbstractVector{<:Real}}

Type for a B-spline basis with breakpoint vector of type `T`.

Here, a B-spline basis is completely specified by its order ``k`` and breakpoint sequence.
The knot sequence is derived from the breakpoint sequence by duplicating the first and last
breakpoints so they each appear ``k`` times. Knot sequences where the first and last
breakpoints do not appear ``k`` times are not supported by this data type.

!!! note
    Here, unlike in some other B-spline libraries, ``k`` always refers to the *order* of a B-spline, not its *degree* (order = degree + 1).
"""
struct BSplineBasis{T<:AbstractVector{<:Real}}
    order::Int
    breakpoints::T

    function BSplineBasis{T}(order, breakpoints) where T<:AbstractVector{<:Real}
        order ≥ 1 || throw(DomainError(order, "order of B-splines must be positive."))
        length(breakpoints) ≥ 2 || throw(ArgumentError("length of breakpoint vector must be at least 2."))
        has_offset_axes(breakpoints) && throw(ArgumentError("breakpoint vector must not have offset axes."))
        new(order, breakpoints)
    end
end

"""
    BSplineBasis(order, breakpoints)

Create a B-spline basis with order `order` and breakpoint vector `breakpoints`. The
breakpoint vector is assumed to be sorted.
"""
BSplineBasis(order, breakpoints) = BSplineBasis{typeof(breakpoints)}(order, breakpoints)

Base.:(==)(x::BSplineBasis, y::BSplineBasis) =
    order(x) == order(y) && breakpoints(x) == breakpoints(y)

Base.hash(x::BSplineBasis, h::UInt) =
    hash(breakpoints(x), hash(order(x), hash(:BSplineBasis, h)))

function Base.checkbounds(b::BSplineBasis, i)
    checkbounds(Bool, b, i) || Base.throw_boundserror(b, i)
    nothing
end
Base.checkbounds(::Type{Bool}, b::BSplineBasis, i) = checkindex(Bool, eachindex(b), i)

Base.eltype(T::Type{<:BSplineBasis}) = BSpline{T}

Base.length(b::BSplineBasis) = Int(length(breakpoints(b))) + order(b) - 2

Base.iterate(b::BSplineBasis, i=1) = i-1 < length(b) ? (@inbounds b[i], i+1) : nothing

Base.firstindex(b::BSplineBasis) = 1
Base.lastindex(b::BSplineBasis)  = length(b)

Base.keys(b::BSplineBasis) = Base.OneTo(lastindex(b))

@propagate_inbounds function Base.getindex(b::BSplineBasis, i::Integer)
    @boundscheck checkbounds(b, i)
    @inbounds BSpline(b, i)
end

Base.iterate(r::Iterators.Reverse{<:BSplineBasis}, i=length(r)) =
    iszero(i) ? nothing : (@inbounds r.itr[i], i-1)

function Base.show(io::IO, ::MIME"text/plain", basis::BSplineBasis)
    summary(io, basis); println(io, ':')
    println(io, " order: ", order(basis))
    print(io, " breakpoints: ")
    show(IOContext(io, :compact=>true), breakpoints(basis))
end

Base.summary(io::IO, basis::BSplineBasis) =
    print(io, length(basis), "-element ", typeof(basis))

"""
    breakpoints(basis::BSplineBasis)

Return the breakpoint sequence of the B-spline basis.

# Examples

```jldoctest
julia> breakpoints(BSplineBasis(3, 0:5))
0:5

julia> breakpoints(BSplineBasis(4, [1.0, 1.5, 2.5, 4.0]))
4-element $(Vector{Float64}):
 1.0
 1.5
 2.5
 4.0
```
"""
breakpoints(b::BSplineBasis) = b.breakpoints

"""
    knots(basis::BSplineBasis)

Return the knot sequence of the B-spline basis.

The knot sequence is the breakpoint sequence except that the first and last values are
duplicated so they appear `order(basis)` times.

# Examples

```jldoctest
julia> knots(BSplineBasis(3, 0:5))
10-element $(BSplines.KnotVector{Int64, UnitRange{Int64}}):
 0
 0
 0
 1
 2
 3
 4
 5
 5
 5
```
"""
knots(b::BSplineBasis) = @inbounds KnotVector(breakpoints(b), order(b)-1)

"""
    order(spline::Spline)
    order(basis::BSplineBasis)

Return the order of a spline or a B-spline basis.

# Examples

```jldoctest
julia> order(BSplineBasis(3, 0:5))
3

julia> order(BSplineBasis(4, [1.0, 1.5, 2.5, 4.0]))
4
```
"""
function order end

order(b::BSplineBasis) = b.order

"""
    intervalindex(vec, x[, start])

If `v` is an `AbstractVector`, return the largest index `i` so that `vec[i] ≤ x ≤ vec[end]`
and `vec[i] < vec[end]`. Return `nothing` if no such index exists. The vector `vec` is
assumed to be sorted in ascending order.

If `vec` is a [`BSplineBasis`](@ref), return `intervalindex(knots(vec), x[, start])`.

If `start` is given, a linear search is performed, starting from the index `start` going
forward or backward. If `start` is not given, a binary search is performed.

# Examples

```jldoctest
julia> intervalindex([1,1,2,3,4,4,4,5,6,6], 2.5)
3

julia> intervalindex([1,1,2,3,4,4,4,5,6,6], 7) # returns nothing

julia> intervalindex([1,1,2,3,4,4,4,5,6,6], 1)
2

julia> intervalindex([1,1,2,3,4,4,4,5,6,6], 4)
7

julia> intervalindex([1,1,2,3,4,4,4,5,6,6], 6.0)
8
```
"""
function intervalindex end

function intervalindex(vec::AbstractVector, x)
    isnan(x) && return nothing
    lastelement = last(vec)
    x < first(vec) && return nothing
    x > lastelement && return nothing
    if x == lastelement
        for i = lastindex(vec)-1:-1:firstindex(vec)
            vec[i] < lastelement && return Int(i)
        end
        return nothing
    end
    return Int(searchsortedlast(vec, x))
end

function intervalindex(vec::AbstractRange, x)
    isnan(x) && return nothing
    firstelement = first(vec)
    lastelement = last(vec)
    x < firstelement && return nothing
    x > lastelement && return nothing
    firstelement == lastelement && return nothing
    x == lastelement && return Int(lastindex(vec)-1)
    return Int(searchsortedlast(vec, x))
end

function intervalindex(vec::AbstractVector, x, start::Integer)
    if vec[start] ≤ x
        ind = findnext(v -> v > x, vec, start+1)
        ind !== nothing && return Int(ind-1)
        lastelement = last(vec)
        if x == lastelement
            for i = lastindex(vec)-1:-1:firstindex(vec)
                vec[i] < lastelement && return Int(i)
            end
        end
        return nothing
    else
        ind = findprev(v -> v ≤ x, vec, start-1)
        ind === nothing && return nothing
        return Int(ind)
    end
end

function intervalindex(vec::AbstractRange, x, start::Integer)
    if vec[start] ≤ x
        ind = findnext(v -> v > x, vec, start+1)
        ind !== nothing && return Int(ind-1)
        lastelement = last(vec)
        first(vec) == lastelement && return nothing
        x == lastelement && return Int(lastindex(vec)-1)
        return nothing
    else
        ind = findprev(v -> v ≤ x, vec, start-1)
        ind === nothing && return nothing
        return Int(ind)
    end
end

intervalindex(b::BSplineBasis, args...) = intervalindex(knots(b), args...)

"""
    check_intervalindex(knots, x, index)

Throw an `ArgumentError` if `index` is not equal to `intervalindex(knots, x)`. (not
exported)
"""
function check_intervalindex(knots, x, index)
    if !check_intervalindex(Bool, knots, x, index)
        throw(ArgumentError("wrong intervalindex: $(repr(index))"))
    end
end

"""
    check_intervalindex(Bool, knots, x, index)

Return `true` if `index` is equal to `intervalindex(knots, x)`. (not exported)
"""
check_intervalindex(::Type{Bool}, knots, x, index)

check_intervalindex(::Type{Bool}, knots::AbstractVector, x, ::Nothing) =
    (x < first(knots)) | (x > last(knots)) | isnan(x)

function check_intervalindex(::Type{Bool}, knots::AbstractVector, x, index::Integer)
    checkbounds(Bool, knots, index)   || return false
    checkbounds(Bool, knots, index+1) || return false
    tl = knots[index]
    tr = knots[index+1]
    (tl ≤ x) & (tl < tr) & ((x < tr) | (x == tr == last(knots)))
end

check_intervalindex(::Type{Bool}, b::BSplineBasis, x, index) =
    check_intervalindex(Bool, knots(b), x, index)

struct IntervalIndices{T<:AbstractVector{<:Real}}
    vec::T
    indices::UnitRange{Int}
    offset::Int
end

"""
    IntervalIndices(vec, indices, offset)

Return an iterator that yields the numbers
`(i + offset for i = indices[1:end-1] if vec[i] < vec[i+1])`. (not exported)
"""
IntervalIndices(vec::T, indices, offset) where T<:AbstractVector{<:Real} =
    IntervalIndices{T}(vec, indices, offset)

Base.IteratorSize(::Type{<:IntervalIndices}) = Base.SizeUnknown()

Base.eltype(::Type{<:IntervalIndices}) = Int

function Base.iterate(i::IntervalIndices, (index,value)=(first(i.indices),i.vec[first(i.indices)]))
    while index < last(i.indices)
        nextindex = index + 1
        nextvalue = i.vec[nextindex]
        value < nextvalue && return (index+i.offset, (nextindex, nextvalue))
        index = nextindex
    end
    return nothing
end

function Base.iterate(i::Iterators.Reverse{<:IntervalIndices})
    isempty(i.itr.indices) && return nothing
    index = last(i.itr.indices)
    iterate(i, (index,i.itr.vec[index]))
end
function Base.iterate(i::Iterators.Reverse{<:IntervalIndices}, (index,value))
    while index > first(i.itr.indices)
        nextindex = index - 1
        nextvalue = i.itr.vec[nextindex]
        value > nextvalue && return (nextindex+i.itr.offset, (nextindex, nextvalue))
        index = nextindex
    end
    return nothing
end

"""
    intervalindices(basis::BSplineBasis, indices=eachindex(basis))

Return an iterator that yields the indices of all intervals on which `basis` is defined,
i.e., it produces all indices `ind` (in ascending order) for which
`(knots(basis)[ind], knots(basis)[ind+1])` is such an interval.

If a range of `indices` is supplied, the iterator yields only those intervals on which *at
least one* of the B-splines `basis[j] for j=indices` is non-zero.

# Examples

```jldoctest
julia> intervalindices(BSplineBasis(3, 0:5))
3:7

julia> intervalindices(BSplineBasis(3, 0:5), 1:4)
3:6

julia> intervalindices(BSplineBasis(4, [1,2,3,4,4,4,5,6]))
$(BSplines.IntervalIndices{Vector{Int64}})([1, 2, 3, 4, 4, 4, 5, 6], 1:8, 3)

julia> collect(ans)
5-element $(Vector{Int64}):
  4
  5
  6
  9
 10
```
"""
intervalindices(::BSplineBasis, ::Union{AbstractUnitRange,Colon}=Colon())

function intervalindices(basis::BSplineBasis, indices::Colon=Colon())
    bps = breakpoints(basis)
    _intervalindices(bps, eachindex(bps), order(basis)-1)
end

function intervalindices(basis::BSplineBasis, indices::AbstractUnitRange)
    checkbounds(basis, indices)
    bps = breakpoints(basis)
    km1 = order(basis)-1
    if isempty(indices)
        _intervalindices(bps, 1:0, km1)
    else
        firstknotindex = max(first(indices)-km1, firstindex(bps))
        lastknotindex  = min(last(indices)+1, lastindex(bps))
        _intervalindices(bps, firstknotindex:lastknotindex, km1)
    end
end

"""
    intervalindices(basis::BSplineBasis, i, j, ...)

For integers `i`, `j`, …, return an iterator that yields the indices of all intervals on
which *all* of the B-splines `basis[i]`, `basis[j]`, … are non-zero, i.e., it produces all
indices `ind` (in ascending order) for which `(knots(basis)[ind], knots(basis)[ind+1])` is
such an interval.

# Examples

```jldoctest
julia> intervalindices(BSplineBasis(3, 0:5), 3)
3:5

julia> intervalindices(BSplineBasis(3, 0:5), 4, 5)
5:6

julia> intervalindices(BSplineBasis(3, 0:5), 2, 6) # B-splines do not overlap
6:5

julia> intervalindices(BSplineBasis(3, 0:5), 3, 5, 4)
5:5

julia> intervalindices(BSplineBasis(4, [1,2,3,4,4,4,5,6]), 3, 5)
$(BSplines.IntervalIndices{Vector{Int64}})([1, 2, 3, 4, 4, 4, 5, 6], 2:4, 3)

julia> collect(ans)
2-element $(Vector{Int64}):
 5
 6
```
"""
function intervalindices(basis::BSplineBasis, indices::Integer...)
    for i in indices
        checkbounds(basis, i)
    end
    bps = breakpoints(basis)
    km1 = order(basis)-1
    firstknotindex = max(max(indices...)-km1, firstindex(bps))
    lastknotindex  = min(min(indices...)+1, lastindex(bps))
    _intervalindices(bps, firstknotindex:lastknotindex, km1)
end

_intervalindices(vec, indices, offset) = IntervalIndices(vec, indices, offset)

_intervalindices(vec::AbstractRange, indices, offset) =
    (iszero(step(vec)) ? indices[1:0] : indices[1:end-1]) .+ offset

length_indices(basis::BSplineBasis, ::Colon) = length(basis)
length_indices(basis::BSplineBasis, indices) = length(indices)

range_indices(basis::BSplineBasis, ::Colon) = eachindex(basis)
range_indices(basis::BSplineBasis, indices) = indices

"""
    support(basis::BSplineBasis) -> a, b

Return the interval ``[a,b]`` on which the B-spline basis is defined, i.e., `a` is the first
and `b` the last breakpoint of the `basis`.

# Examples

```jldoctest
julia> support(BSplineBasis(3, 0:5))
(0, 5)

julia> support(BSplineBasis(4, [1.0, 1.5, 2.5, 4.0]))
(1.0, 4.0)
```
"""
support(b::BSplineBasis) = (bps = breakpoints(b); (first(bps), last(bps)))

"""
    bsplines(basis, x[, derivative]; kwargs...)

Calculate the values and/or derivatives of all non-zero B-splines of `basis` at `x`.

If all B-splines of `basis` are zero at `x`, `nothing` is returned. If any B-splines are
non-zero at `x`, an `OffsetArray` is returned. Its size and contents depend on the optional
`derivative` argument:
* If `derivative` is not given, the returned `OffsetArray` contains the value of the `i`-th
  B-spline at the index `i`.
* If `derivative == Derivative(N)`, the returned `OffsetArray` contains the `N`-th
  derivative of the `i`-th B-spline at the index `i`.
* If `derivative == AllDerivatives(N)`, the returned `OffsetArray` contains the `m`-th
  derivative (`0 ≤ m < N`) of the `i`-th B-spline at the index `i, m`.

Two optional keyword arguments can be used to increase performance:
* `derivspace`: When calculating derivatives, some coefficients are stored in a matrix of
  size `(order(basis), order(basis))`. By default, the function allocates a new matrix. To
  avoid this, a pre-allocated matrix can be supplied with the `derivspace` keyword. It can
  only be used when calculating derivatives, i.e., with `Derivative(N)` where `N ≥ 1` or
  `AllDerivatives(N)` where `N ≥ 2`. To also pre-allocate the array that contains the
  result, use the [`bsplines!`](@ref) function instead.
* `leftknot`: If the index of the relevant interval (i.e.,
  `intervalindex(basis(spline), x)`) is already known, it can be supplied with the
  `leftknot` keyword.

# Examples

```jldoctest
julia> basis = BSplineBasis(4, 0:5);

julia> bsplines(basis, 2.4)
4-element OffsetArray(::$(Vector{Float64}), 3:6) with eltype Float64 with indices 3:6:
 0.03600000000000002
 0.5386666666666667
 0.41466666666666663
 0.01066666666666666

julia> bsplines(basis, 2.4, Derivative(1), derivspace=zeros(4,4))
4-element OffsetArray(::$(Vector{Float64}), 3:6) with eltype Float64 with indices 3:6:
 -0.18000000000000005
 -0.5599999999999999
  0.66
  0.07999999999999996

julia> bsplines(basis, 6) # returns nothing

julia> bsplines(basis, 17//5, leftknot=7)
4-element OffsetArray(::$(Vector{Rational{Int64}}), 4:7) with eltype Rational{Int64} with indices 4:7:
   9//250
 202//375
 307//750
   2//125

julia> bsplines(basis, 2.4, AllDerivatives(3))
4×3 OffsetArray(::$(Matrix{Float64}), 3:6, 0:2) with eltype Float64 with indices 3:6×0:2:
 0.036      -0.18   0.6
 0.538667   -0.56  -0.8
 0.414667    0.66  -0.2
 0.0106667   0.08   0.4
```
"""
function bsplines(basis::BSplineBasis, x, drv=NoDerivative();
                  leftknot=intervalindex(basis, x), derivspace=nothing)
    check_intervalindex(basis, x, leftknot)
    check_derivspace(basis, drv, derivspace)
    leftknot === nothing && return nothing
    dest = bsplines_destarray(basis, x, drv, derivspace)
    offset = @inbounds _bsplines!(dest, derivspace, basis, x, leftknot, drv)
    bsplines_offsetarray(dest, offset, drv)
end

"""
    bsplines!(dest, basis, x[, derivative]; kwargs...)

Calculate the values and/or derivatives of all non-zero B-splines of `basis` at `x` in-place
(i.e., in `dest`).

If all B-splines of `basis` are zero at `x`, `nothing` is returned. If any B-splines are
non-zero at `x`, an `OffsetArray` that wraps `dest` is returned. Its size and contents
depend on the optional `derivative` argument:
* If `derivative` is not given, the returned `OffsetArray` contains the value of the `i`-th
  B-spline at the index `i`. In this case, `dest` must be a vector of length `order(basis)`.
* If `derivative == Derivative(N)`, the returned `OffsetArray` contains the `N`-th
  derivative of the `i`-th B-spline at the index `i`. Again, `dest` must be a vector of
  length `order(basis)`.
* If `derivative == AllDerivatives(N)`, the returned `OffsetArray` contains the `m`-th
  derivative (`0 ≤ m < N`) of the `i`-th B-spline at the index `i, m`. In this case, `dest`
  must be a matrix of size `(order(basis), N)`.

Two optional keyword arguments can be used to increase performance:
* `derivspace`: When calculating derivatives, some coefficients are stored in a matrix of
  size `(order(basis), order(basis))`. By default, the function allocates a new matrix. To
  avoid this, a pre-allocated matrix can be supplied with the `derivspace` keyword. It can
  only be used when calculating derivatives, i.e., with `Derivative(N)` where `N ≥ 1` or
  `AllDerivatives(N)` where `N ≥ 2`.
* `leftknot`: If the index of the relevant interval (i.e.,
  `intervalindex(basis(spline), x)`) is already known, it can be supplied with the
  `leftknot` keyword.

# Examples

```jldoctest
julia> basis = BSplineBasis(4, 0:5);

julia> dest = zeros(4);

julia> bsplines!(dest, basis, 2.4)
4-element OffsetArray(::$(Vector{Float64}), 3:6) with eltype Float64 with indices 3:6:
 0.03600000000000002
 0.5386666666666667
 0.41466666666666663
 0.01066666666666666

julia> parent(ans) === dest
true

julia> bsplines!(dest, basis, -1.0) # returns nothing

julia> dest = zeros(4, 3);

julia> bsplines!(dest, basis, 3.75, AllDerivatives(3))
4×3 OffsetArray(::$(Matrix{Float64}), 4:7, 0:2) with eltype Float64 with indices 4:7×0:2:
 0.00260417  -0.03125    0.25
 0.315104    -0.65625    0.25
 0.576823     0.265625  -1.625
 0.105469     0.421875   1.125

julia> parent(ans) === dest
true
```
"""
function bsplines!(dest, basis::BSplineBasis, x, drv=NoDerivative();
                   leftknot=intervalindex(basis, x), derivspace=nothing)
    check_intervalindex(basis, x, leftknot)
    check_destarray_axes(dest, basis, drv)
    check_derivspace(basis, drv, derivspace)
    leftknot === nothing && return nothing
    offset = @inbounds _bsplines!(dest, derivspace, basis, x, leftknot, drv)
    bsplines_offsetarray(dest, offset, drv)
end

check_derivspace(basis, drv::NoDerivUnion, ::Nothing) = nothing
check_derivspace(basis, drv::NoDerivUnion, derivspace) =
    throw(ArgumentError("derivspace keyword cannot be used with $drv."))

check_derivspace(basis, drv, ::Nothing) = nothing
function check_derivspace(basis, drv, derivspace)
    got = axes(derivspace)
    exp = (Base.OneTo(order(basis)), Base.OneTo(order(basis)))
    if got != exp
        throw(DimensionMismatch("derivspace array has wrong axes: expected $exp, got $got"))
    end
    nothing
end

bsplines_destarray(basis, x, drv, derivspace) = bsplines_destarray(eltype(derivspace), basis, drv)
bsplines_destarray(basis, x, drv, ::Nothing)  = bsplines_destarray(bspline_returntype(basis, x), basis, drv)

bsplines_destarray(T, basis, ::Derivative)                = Vector{T}(undef, order(basis))
bsplines_destarray(T, basis, ::AllDerivatives{N}) where N = Matrix{T}(undef, order(basis), N)

bsplines_offsetarray(arr, offset, ::Derivative)     = OffsetArray(arr, offset)
bsplines_offsetarray(arr, offset, ::AllDerivatives) = OffsetArray(arr, offset, -1)

function check_destarray_axes(dest, basis, drv)
    got = axes(dest)
    exp = destarray_axes(basis, drv)
    if got != exp
        throw(DimensionMismatch("destination array has wrong axes: expected $exp, got $got"))
    end
    nothing
end

destarray_axes(basis, ::Derivative)                = (Base.OneTo(order(basis)),)
destarray_axes(basis, ::AllDerivatives{N}) where N = (Base.OneTo(order(basis)), Base.OneTo(N))

# The implementation of this method (and the method iterate_bsplines! it calls) is adapted
# from the Fortran subroutine BSPLVB from Carl de Boor’s book *A practical Guide to
# Splines* [^deBoor1978].
#
# [^deBoor1978]:
#     Carl de Boor, *A Practical Guide to Splines*, New York, N.Y.: Springer-Verlag, 1978.
@propagate_inbounds function _bsplines!(dest, _, basis, x, leftknot::Integer, ::NoDerivUnion)
    t = knots(basis)
    k = order(basis)
    xtyped = convert(eltype(dest), x)
    dest[1] = oneunit(eltype(dest))
    for j = 1:k-1
        iterate_bsplines!(dest, dest, t, j, xtyped, leftknot)
    end
    return leftknot-k
end

# The implementation of this method (and the methods iterate_bsplines! and
# iterate_derivatives! it calls) is adapted from the Fortran subroutine BSPLVD from Carl de
# Boor’s book *A practical Guide to Splines* [^deBoor1978].
#
# [^deBoor1978]:
#     Carl de Boor, *A Practical Guide to Splines*, New York, N.Y.: Springer-Verlag, 1978.
@propagate_inbounds function _bsplines!(dest, derivspace, basis, x, leftknot::Integer, ::AllDerivatives{N}) where N
    t = knots(basis)
    k = order(basis)
    xtyped = convert(eltype(dest), x)
    N′ = min(k, N)
    # Set k-th and higher derivatives to zero
    dest[:, N′+1:N] .= zero(eltype(dest))
    # Successively calculate values of B-splines and store them in dest
    lastcol = @view dest[N′:k, N′]
    lastcol[1] = oneunit(eltype(dest))
    for j = 1:k-N′
        iterate_bsplines!(lastcol, lastcol, t, j, xtyped, leftknot)
    end
    newcol = lastcol
    for j = k-N′+1:k-1
        oldcol = newcol
        newcol = @view dest[k-j:k, k-j]
        iterate_bsplines!(newcol, oldcol, t, j, xtyped, leftknot)
    end
    # Successively calculate coefficients for derivatives and combine them with B-splines
    drvcoeffs = prepare_drvcoeffs(eltype(dest), order(basis), derivspace)
    for col = 2:N′
        # Column `col` contains `col-1`-th derivatives
        iterate_derivatives!(drvcoeffs, t, k, col-1, leftknot)
        for row = 1:k
            sum = zero(eltype(dest))
            for j = max(row,col):k
                sum += oftype(sum, drvcoeffs[j, row] * dest[j, col])
            end
            dest[row, col] = sum
        end
    end
    return leftknot-k
end

prepare_drvcoeffs(T, k, ::Nothing) = Matrix{T}(I, k, k)
function prepare_drvcoeffs(_, k, derivspace)
    fill!(derivspace, 0)
    for i = diagind(derivspace)
        derivspace[i] = 1
    end
    derivspace
end

# The implementation of this method (and the methods iterate_bsplines! and
# iterate_derivatives! it calls) is adapted from the Fortran subroutine BSPLVD from Carl de
# Boor’s book *A practical Guide to Splines* [^deBoor1978].
#
# [^deBoor1978]:
#     Carl de Boor, *A Practical Guide to Splines*, New York, N.Y.: Springer-Verlag, 1978.
@propagate_inbounds function _bsplines!(dest, derivspace, basis, x, leftknot::Integer, ::Derivative{N}) where N
    t = knots(basis)
    k = order(basis)
    xtyped = convert(eltype(dest), x)
    if N ≥ k
        # k-th and higher derivatives are zero
        dest .= zero(eltype(dest))
    else
        # Calculate B-splines of order k-N and store them in dest[N+1:k]
        col = @view dest[N+1:k]
        # col = uview(dest, N+1:k)
        col[1] = oneunit(eltype(dest))
        for j = 1:k-N-1
            iterate_bsplines!(col, col, t, j, xtyped, leftknot)
        end
        # Calculate coefficients for derivatives and combine them with B-splines
        drvcoeffs = prepare_drvcoeffs(eltype(dest), order(basis), derivspace)
        for drv = 1:N
            iterate_derivatives!(drvcoeffs, t, k, drv, leftknot)
        end
        # Combine coefficients with B-splines
        for i = 1:k
            sum = zero(eltype(dest))
            for j = max(i,N+1):k
                sum += oftype(sum, drvcoeffs[j,i] * dest[j])
            end
            dest[i] = sum
        end
    end
    return leftknot-k
end

# Take the B-splines of order `j` from `src[1:j]` and write the B-splines of order `j+1` to `dest[1:j+1]`
@propagate_inbounds function iterate_bsplines!(dest::AbstractArray{T}, src::AbstractArray{T}, knots, j, x::T, leftknot) where T
    saved = zero(T)
    for r = 1:j
        tₗ₊ᵣ   = convert(T, knots[leftknot+r])
        tₗ₊ᵣ₋ⱼ = convert(T, knots[leftknot+r-j])
        term = src[r] / (tₗ₊ᵣ - tₗ₊ᵣ₋ⱼ)
        dest[r] = saved + (tₗ₊ᵣ - x) * term
        saved = oftype(saved, (x - tₗ₊ᵣ₋ⱼ) * term)
    end
    dest[j+1] = saved
end

# Iterate the coefficients in `coeffs` from the `deriv-1`-th to the `deriv`-th derivative
@propagate_inbounds function iterate_derivatives!(coeffs::AbstractArray{T}, knots, k, deriv, leftknot) where T
    kmd = k-deriv
    for j = 1:kmd
        lp1mj = leftknot+1-j
        factor = kmd / (convert(T, knots[lp1mj+kmd]) - convert(T, knots[lp1mj]))
        row = k+1-j
        for col = 1:row
            coeffs[row, col] = (coeffs[row, col] - coeffs[row-1, col]) * factor
        end
    end
end

"""
    basismatrix(basis::BSplineBasis, xvalues; indices=eachindex(basis))

Calculate the matrix `[basis[i](x) for x=xvalues, i=indices]`.

# Examples

```jldoctest
julia> basis = BSplineBasis(3, 0:5);

julia> x = [0.3, 1.5, 3.2, 4.5];

julia> basismatrix(basis, x)
4×7 $(Matrix{Float64}):
 0.49  0.465  0.045  0.0    0.0    0.0    0.0
 0.0   0.125  0.75   0.125  0.0    0.0    0.0
 0.0   0.0    0.0    0.32   0.66   0.02   0.0
 0.0   0.0    0.0    0.0    0.125  0.625  0.25
```
"""
function basismatrix(basis::BSplineBasis, xvalues; indices=Colon())
    checkbounds(Bool, basis, indices) || throw(ArgumentError("invalid indices for basis: $indices"))
    T = bspline_returntype(basis, eltype(xvalues))
    dest = zeros(T, length(xvalues), length_indices(basis, indices))
    derivspace = similar(dest, order(basis))
    _basismatrix!(dest, derivspace, basis, xvalues, indices)
    dest
end

"""
    basismatrix!(dest, basis::BSplineBasis, xvalues; indices=eachindex(basis))

Calculate the matrix `[basis[i](x) for x=xvalues, i=indices]` and store it in `dest`.

# Examples

```jldoctest
julia> basis = BSplineBasis(3, 0:5);

julia> x = [0.3, 1.5, 3.2, 4.5];

julia> dest = Array{Float64}(undef, length(x), length(basis));

julia> basismatrix!(dest, basis, x)
4×7 $(Matrix{Float64}):
 0.49  0.465  0.045  0.0    0.0    0.0    0.0
 0.0   0.125  0.75   0.125  0.0    0.0    0.0
 0.0   0.0    0.0    0.32   0.66   0.02   0.0
 0.0   0.0    0.0    0.0    0.125  0.625  0.25
```
"""
function basismatrix!(dest, basis::BSplineBasis, xvalues; indices=Colon())
    has_offset_axes(dest) && throw(ArgumentError("destination matrix may not have offset axes."))
    checkbounds(Bool, basis, indices) || throw(ArgumentError("invalid indices for basis: $indices"))
    expected_size = (length(xvalues), length_indices(basis, indices))
    size(dest) == expected_size || throw(DimensionMismatch("destination array has size $(size(dest)), expected $expected_size"))
    workspace = similar(dest, order(basis))
    fill!(dest, zero(eltype(dest)))
    _basismatrix!(dest, workspace, basis, xvalues, indices)
    dest
end

function _basismatrix!(dest, workspace, basis::BSplineBasis, xvalues, indices::AbstractUnitRange)
    indicesoffset = first(indices) - 1
    start = 1
    for (xindex, x) = enumerate(xvalues)
        leftknot = intervalindex(basis, x, start)
        @assert leftknot !== nothing "xvalue outside of the support of basis: $x."
        start = leftknot
        bspl = bsplines!(workspace, basis, x, leftknot=leftknot)
        for index = axes(bspl, 1) ∩ indices
            dest[xindex, index-indicesoffset] = bspl[index]
        end
    end
end

function _basismatrix!(dest, workspace, basis::BSplineBasis, xvalues, ::Colon)
    start = 1
    for (xindex, x) = enumerate(xvalues)
        leftknot = intervalindex(basis, x, start)
        @assert leftknot !== nothing "xvalue outside of the support of basis: $x."
        start = leftknot
        bspl = bsplines!(workspace, basis, x, leftknot=leftknot)
        for index = axes(bspl, 1)
            dest[xindex, index] = bspl[index]
        end
    end
end


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %% [markdown]
# https://twitter.com/physics303/status/1414522243216859143?s=21

# %%
f(A) = sortperm(sort.(collect(eachrow(A))); rev=true)

function g!(A, ix, tmp)
    for i in axes(A, 1) tmp[i] .= @view A[i, :] end
    sort!.(tmp)
    sortperm!(ix, tmp; rev=true)
end

function maketmp(A)
    ix = Vector{Int}(undef, size(A, 1))
    tmp = [Vector{eltype(A)}(undef, size(A, 2)) for _ in 1:size(A, 1)]
    ix, tmp
end

B = [
    2 4 4 4 5 9
    1 2 2 3 5 9
    1 2 3 3 3 9
    1 2 2 5 6 9
    1 3 4 4 4 9
]

ix, tmp = maketmp(B)
@show ix
@show tmp
@show f(B) == g!(B, ix, tmp)
@time f(B)
@time g!(B, ix, tmp)

# %%
using BenchmarkTools
@btime f($B)
@btime g!($B, $ix, $tmp)

# %%


function varargout = regression(varargin)
#  REGRESSION MATLAB code for regression.fig
#       REGRESSION, by itself, creates a new REGRESSION or raises the existing
#       singleton*.
#
#       H = REGRESSION returns the handle to a new REGRESSION or the handle to
#       the existing singleton*.
#
#       REGRESSION('CALLBACK',hObject,eventData,handles,...) calls the local
#       function named CALLBACK in REGRESSION.M with the given input arguments.
#
#       REGRESSION('Property','Value',...) creates a new REGRESSION or raises the
#       existing singleton*.  Starting from the left, property value pairs are
#       applied to the GUI before regression_OpeningFcn gets called.  An
#       unrecognized property name or invalid value makes property application
#       stop.  All inputs are passed to regression_OpeningFcn via varargin.
#
#       *See GUI Options on GUIDE's Tools menu.  Choose "GUI allows only one
#       instance to run (singleton)".
#
#  See also: GUIDE, GUIDATA, GUIHANDLES

#  Edit the above text to modify the response to help regression

#  Last Modified by GUIDE v2.5 17-Apr-2013 12:46:04

#  Begin initialization code - DO NOT EDIT
gui_Singleton = 1;
gui_State = struct('gui_Name',       mfilename, ...
                   'gui_Singleton',  gui_Singleton, ...
                   'gui_OpeningFcn', @regression_OpeningFcn, ...
                   'gui_OutputFcn',  @regression_OutputFcn, ...
                   'gui_LayoutFcn',  [] , ...
                   'gui_Callback',   []);
if nargin && ischar(varargin{1})
    gui_State.gui_Callback = str2func(varargin{1});
end

if nargout
    [varargout{1:nargout}] = gui_mainfcn(gui_State, varargin{:});
else
    gui_mainfcn(gui_State, varargin{:});
end
#  End initialization code - DO NOT EDIT

#  --- Executes just before regression is made visible.
function regression_OpeningFcn(hObject, ~, handles, varargin)
#  This function has no output args, see OutputFcn.
#  hObject    handle to figure
#  eventdata  reserved - to be defined in a future version of MATLAB
#  handles    structure with handles and user data (see GUIDATA)
#  varargin   command line arguments to regression (see VARARGIN)

#  Choose default command line output for regression
handles.output = hObject;

#  Update handles structure
guidata(hObject, handles);

#  UIWAIT makes regression wait for user response (see UIRESUME)
#  uiwait(handles.figure1);

# File: regression.jl
# Perform a linear regression
# Author:
# Date: 29-June-2020

# import XL data
[~,~,returns] = getXLData('hfma_matlab_data.xlsx','LS Index & S&P 500');
# S&P 500 (# )
x = returns(:,2);
# Long-Short Index (# )
y = returns(:,1);

# requested statistics
list = {'beta','adjrsquare','rsquare','tstat'};
# perform regression
stats = regstats(y,x,'linear',list);

# output statistics
cfs = stats.beta; # regression coefficients   β  \beta
arsq = stats.adjrsquare; # adjusted R square
rsq = stats.rsquare; # R square
t = stats.tstat.t; # t-statistic
pv = stats.tstat.pval; # p-value

# other statistics
se = cfs./t; # standard error
correl = corr(x,y); # correlation

# OUTPUT
# format data 1
data1 = num2cell([correl;arsq;rsq]);
    for i = 1:numel(data1)
            data1{i} = sprintf('# 12.4f',data1{i});
    end
# set column headings
headings1 = {'Correlation','adjusted R2','R2'};
# set table output
set(handles.table_regression_results_1,'ColumnName','','RowName',headings1,'Data',data1);
# format data 2
data2 = num2cell([cfs,se,t,pv]);
    for i = 1:numel(data2)
        data2{i} = sprintf('# 12.4f',data2{i});
    end
# set size column widths
cwidth = num2cell(repmat(70,1,10));
# set column headings
headings2 = {'Coefficients','SE','t-statistic','p-Value'};
# set table output
set(handles.table_regression_results_2,'ColumnName',headings2,'ColumnWidth',cwidth,'RowName',{'Intercept (a)','Slope (b) [S&P500]'},'Data',data2);

# plot linear regression
p = polyfit(x,y,1);
f = polyval(p,x);

figure()
plot(x,y,'.');
hold on
# linear fit
plot(x,f,'-r');
xlabel('S&P500 (# )');
ylabel('Long-Short Index (# )');
# add linear equation & R-squared
text(-1,-5,['y= ',num2str(cfs(2)),'x + ',num2str(cfs(1)),',  R^2= ',num2str(rsq)]);

#  --- Outputs from this function are returned to the command line.
function varargout = regression_OutputFcn(~, ~, handles)
#  varargout  cell array for returning output args (see VARARGOUT);
#  hObject    handle to figure
#  eventdata  reserved - to be defined in a future version of MATLAB
#  handles    structure with handles and user data (see GUIDATA)

#  Get default command line output from handles structure
varargout{1} = handles.output;


# ---
# title: 1406. Stone Game III
# id: problem1406
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/stone-game-iii/description/>
# hidden: true
# ---
#
# Alice and Bob continue their games with piles of stones. There are several
# stones  **arranged in a row** , and each stone has an associated value which
# is an integer given in the array `stoneValue`.
#
# Alice and Bob take turns, with **Alice** starting first. On each player's
# turn, that player can take **1, 2 or 3 stones**  from the **first** remaining
# stones in the row.
#
# The score of each player is the sum of values of the stones taken. The score
# of each player is **0**  initially.
#
# The objective of the game is to end with the highest score, and the winner is
# the player with the highest score and there could be a tie. The game continues
# until all the stones have been taken.
#
# Assume Alice and Bob  **play optimally**.
#
# Return _" Alice"_ if Alice will win, _" Bob"_ if Bob will win or _" Tie"_ if
# they end the game with the same score.
#
#
#
# **Example 1:**
#
#
#
#     Input: values = [1,2,3,7]
#     Output: "Bob"
#     Explanation: Alice will always lose. Her best move will be to take three piles and the score become 6. Now the score of Bob is 7 and Bob wins.
#
#
# **Example 2:**
#
#
#
#     Input: values = [1,2,3,-9]
#     Output: "Alice"
#     Explanation: Alice must choose all the three piles at the first move to win and leave Bob with negative score.
#     If Alice chooses one pile her score will be 1 and the next move Bob's score becomes 5. The next move Alice will take the pile with value = -9 and lose.
#     If Alice chooses two piles her score will be 3 and the next move Bob's score becomes 3. The next move Alice will take the pile with value = -9 and also lose.
#     Remember that both play optimally so here Alice will choose the scenario that makes her win.
#
#
# **Example 3:**
#
#
#
#     Input: values = [1,2,3,6]
#     Output: "Tie"
#     Explanation: Alice cannot win this game. She can end the game in a draw if she decided to choose all the first three piles, otherwise she will lose.
#
#
# **Example 4:**
#
#
#
#     Input: values = [1,2,3,-1,-2,-3,7]
#     Output: "Alice"
#
#
# **Example 5:**
#
#
#
#     Input: values = [-1,-2,-3]
#     Output: "Tie"
#
#
#
#
# **Constraints:**
#
#   * `1 <= values.length <= 50000`
#   * `-1000 <= values[i] <= 1000`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end