# # Calculus on the reference triangle
# In this example, we sample a bivariate function:
# ```math
# f(x,y) = \frac{1}{1+x^2+y^2},
# ```
# on the reference triangle with vertices $(0,0)$, $(0,1)$, and $(1,0)$ and analyze it
# in a Proriol series. Then, we find Proriol series for each component of its
# gradient by term-by-term differentiation of our expansion, and we compare them
# with the true Proriol series by sampling an exact expression for the gradient.
#
# We analyze $f(x,y)$ on an $N\times M$ mapped tensor product grid defined by:
# ```math
# \begin{aligned}
# x & = (1+u)/2,\quad{\rm and}\quad y = (1-u)(1+v)/4,\quad {\rm where:}\\
# u_n & = \cos\left[(n+\tfrac{1}{2})\pi/N\right],\quad{\rm for}\quad 0\le n < N,\quad{\rm and}\\
# v_m & = \cos\left[(m+\tfrac{1}{2})\pi/M\right],\quad{\rm for}\quad 0\le m < M;
# \end{aligned}
# ```
# we convert the function samples to mapped Chebyshev² coefficients using
# `plan_tri_analysis`; and finally, we transform the mapped Chebyshev²
# coefficients to Proriol coefficients using `plan_tri2cheb`.
#
# For the storage pattern of the arrays, please consult the
# [documentation](https://MikaelSlevinsky.github.io/FastTransforms).

using FastTransforms, LinearAlgebra, Plots
const GENFIGS = joinpath(pkgdir(FastTransforms), "docs/src/generated")
!isdir(GENFIGS) && mkdir(GENFIGS)
plotlyjs()

# Our function $f$ and the Cartesian components of its gradient:
f = (x,y) -> 1/(1+x^2+y^2)
fx = (x,y) -> -2x/(1+x^2+y^2)^2
fy = (x,y) -> -2y/(1+x^2+y^2)^2

# The polynomial degree:
N = 15
M = N

# The parameters of the Proriol series:
α, β, γ = 0, 0, 0

# The $u$ grid:
u = [sinpi((N-2n-1)/(2N)) for n in 0:N-1]

# And the $v$ grid:
v = [sinpi((M-2m-1)/(2M)) for m in 0:M-1]

# Instead of using the $u\times v$ grid, we use one with more accuracy near the origin.
# Defining $x$ by:
x = [sinpi((2N-2n-1)/(4N))^2 for n in 0:N-1]

# And $w$ by:
w = [sinpi((2M-2m-1)/(4M))^2 for m in 0:M-1]

# We see how the two grids are related:
((1 .+ u)./2 ≈ x) * ((1 .- u).*(1 .+ v')/4 ≈ reverse(x).*w')

# On the mapped tensor product grid, our function samples are:
F = [f(x[n+1], x[N-n]*w[m+1]) for n in 0:N-1, m in 0:M-1]

# We superpose a surface plot of $f$ on top of the grid:
X = [x for x in x, w in w]
Y = [x[N-n]*w[m+1] for n in 0:N-1, m in 0:M-1]
scatter3d(vec(X), vec(Y), vec(0F); markersize=0.75, markercolor=:blue)
surface!(X, Y, F; legend=false, xlabel="x", ylabel="y", zlabel="f")
savefig(joinpath(GENFIGS, "proriol.html"))
###```@raw html
###<object type="text/html" data="../proriol.html" style="width:100%;height:400px;"></object>
###```

# We precompute a Proriol--Chebyshev² plan:
P = plan_tri2cheb(F, α, β, γ)

# And an FFTW Chebyshev² plan on the triangle:
PA = plan_tri_analysis(F)

# Its Proriol-$(α,β,γ)$ coefficients are:
U = P\(PA*F)

# Similarly, our function's gradient samples are:
Fx = [fx(x[n+1], x[N-n]*w[m+1]) for n in 0:N-1, m in 0:M-1]

# and:
Fy = [fy(x[n+1], x[N-n]*w[m+1]) for n in 0:N-1, m in 0:M-1]

# For the partial derivative with respect to $x$, [Olver et al.](https://doi.org/10.1137/19M1245888)
# derive simple expressions for the representation of this component
# using a Proriol-$(α+1,β,γ+1)$ series.
Gx = zeros(Float64, N, M)
for m = 0:M-2
    for n = 0:N-2
        cf1 = m == 0 ? sqrt((n+1)*(n+2m+α+β+γ+3)/(2m+β+γ+2)*(m+γ+1)*8) : sqrt((n+1)*(n+2m+α+β+γ+3)/(2m+β+γ+1)*(m+β+γ+1)/(2m+β+γ+2)*(m+γ+1)*8)
        cf2 = sqrt((n+α+1)*(m+1)/(2m+β+γ+2)*(m+β+1)/(2m+β+γ+3)*(n+2m+β+γ+3)*8)
        Gx[n+1, m+1] = cf1*U[n+2, m+1] + cf2*U[n+1, m+2]
    end
end
Px = plan_tri2cheb(Fx, α+1, β, γ+1)
Ux = Px\(PA*Fx)

# For the partial derivative with respect to y, the analogous formulae result
# in a Proriol-$(α,β+1,γ+1)$ series.
Gy = zeros(Float64, N, M)
for m = 0:M-2
    for n = 0:N-2
        Gy[n+1, m+1] = 4*sqrt((m+1)*(m+β+γ+2))*U[n+1, m+2]
    end
end
Py = plan_tri2cheb(Fy, α, β+1, γ+1)
Uy = Py\(PA*Fy)

# The $2$-norm relative error in differentiating the Proriol series
# for $f(x,y)$ term-by-term and its sampled gradient is:
hypot(norm(Ux-Gx), norm(Uy-Gy))/hypot(norm(Ux), norm(Uy))

# This error can be improved upon by increasing $N$ and $M$.


############################
## grossMomentsLogMoments ##
############################

function grossRetMomentsToLogRetMoments(mu::Float64, sigma::Float64)
    ## arithmetic moments R to r^{log}

    ## calculate log variance
    varLog = log(1 + sigma^2/mu^2)
    muLog = log(mu) - 0.5*varLog
    sigmaLog = sqrt(varLog)

    return muLog, sigmaLog
end

function logRetMomentsToGrossRetMoments(muLog::Float64,
                                        sigmaLog::Float64)
    ## logarithmic moments r^{log} to R
    mu = exp(muLog + 0.5*sigmaLog^2)
    varGross = (exp(sigmaLog^2) - 1)*exp(2*muLog + sigmaLog^2)

    return mu, sqrt(varGross)
end

#######################
## scaling functions ##
#######################

## Note: initialized without any scaling by now
## Problem: mus should not be scaled linearly, as discrete returns may
## not be aggregated over time linearly. A better way might be to make
## a distributional assumption (log-normal) and scale geometric means
## (or arithmetic means for log returns)

function defaultMuSigmaScaling(mu::Float64, sigma::Float64;
                               scalingFactor = 52)
    ## - annualized arithmetic mean and sigma shown
    ## - via annualized log moments
    ## - percentage values
    ## Assumptions:
    ## - log-normality to get log moments
    ## - independence over time for square-root-of-time scaling
    ## - constant weights over time
    ## - constant moments over time

    ## net to gross
    muGross = mu + 1

    ## get log moments
    muLog, sigmaLog = grossRetMomentsToLogRetMoments(muGross, sigma)

    ## annualize log moments
    muLogAnnual = muLog*scalingFactor
    sigmaLogAnnual = sqrt(scalingFactor)*sigmaLog

    ## translate to gross moments
    muGrossAnnual, sigmaAnnual =
        logRetMomentsToGrossRetMoments(muLogAnnual,
                                       sigmaLogAnnual)

    ## translate to net returns
    muNetAnnual = muGrossAnnual - 1

    ## translate to percentages
    muNetAnnualPercent = muNetAnnual * 100
    sigmaAnnualPercent = sigmaAnnual * 100

    return muNetAnnualPercent, sigmaAnnualPercent
end

## applied to vectors
function defaultMuSigmaScaling(mus::Array{Float64, 1},
                               sigmas::Array{Float64, 1};
                               scalingFactor = 52)
    ## preallocation
    nMoments = size(mus, 1)
    scaledMus = ones(nMoments)
    scaledSigmas = ones(nMoments)

    for ii=1:nMoments
        thisMu, thisSigma =
            defaultMuSigmaScaling(mus[ii], sigmas[ii])
        scaledMus[ii] = thisMu
        scaledSigmas[ii] = thisSigma
    end
    return scaledMus, scaledSigmas
end

##########################################
## symbol / string conversion functions ##
##########################################

function symbToStr(symbs::Array{Symbol, 1})
    return UTF8String[string(symb) for symb in symbs]
end

function symbToStr(symbs::DataArray{Symbol, 1})
    return UTF8String[string(symb) for symb in symbs]
end

function strToSymb(strs::Array{UTF8String, 1})
    return Symbol[symbol(xx) for xx in strs]
end

function strToSymb(strs::DataArray{UTF8String, 1})
    return Symbol[symbol(xx) for xx in strs]
end

function strToSymb(strs::DataArray{String, 1})
    return Symbol[symbol(xx) for xx in strs]
end

#######################
## calculate moments ##
#######################

## single portfolio, single universe
function getPMean(wgts::Array{Float64, 1}, mus::Array{Float64, 1})
    return (wgts'*mus)[1]
end

## multiple portfolios, single universe
function getPMean(wgts::Array{Float64, 2}, mus::Array{Float64, 1})
    nPfs = size(wgts, 1)
    pfMus = ones(nPfs)
    for ii=1:nPfs
        pfMus[ii] = getPMean(wgts[ii, :][:], mus)
    end
    return pfMus
end

## single portfolio, single universe
function getPVar(wgts::Array{Float64, 1}, covMatr::Array{Float64, 2})
    return (wgts'*covMatr*wgts)[1]
end

## multiple portfolios, single universe
function getPVar(wgts::Array{Float64, 2}, covMatr::Array{Float64, 2})
    nPfs = size(wgts, 1)
    pfVariances = ones(nPfs)
    for ii=1:nPfs
        pfVariances[ii] = getPVar(wgts[ii, :][:], covMatr)
    end
    return pfVariances
end

####################################
## statistic functions DataFrames ##
####################################

import Base.mean
function mean(df::DataFrame, int::Integer)
    if int==2
        error("for row-wise mean, use rowmean instead")
    else
        mus = DataFrame(mean(array(df), 1), names(df))
    end
end

import Base.cov
function cov(df::DataFrame)
    return DataFrame(cov(array(df)), names(df))
end

function corrcov(covMatr::DataFrame)
    d = diagm(1/sqrt(diag(array(covMatr))))
    return DataFrame(d*array(covMatr)*d, names(covMatr))
end

##################################
## create random weights matrix ##
##################################

## function randWgts(nObs::Int, nAss::Int)
##     ## create random weights
##     ##
##     ## Output:
##     ## 	nObs x nAss Array{Float64,2} containing portfolio weights in
##     ## 	each row

##     simVals = rand(nObs, nAss)
##     rowsums = sum(simVals, 2)
##     wgts = simVals ./ repmat(rowsums, 1, nAss)
## end

##################################################
## normalize matrix values to represent weights ##
##################################################

function makeWeights(matr::Array{Float64, 2})
    ## normalize matrix to have row sums of 1
    nObs = size(matr, 1)
    for ii=1:nObs
        matr[ii, :] = matr[ii, :]./sum(matr[ii, :])
    end
    return matr
end

#################################
## download interest rate data ##
#################################

function getTBill()
    ## get Federal funds effective rate

    tbillsURL =
        "http://www.federalreserve.gov/datadownload/Output.aspx?rel=H15&series=646250c87b1afd04cc6774796fc0cec8&lastObs=&from=&to=&filetype=csv&label=include&layout=seriescolumn"

    rawSplitted = download(tbillsURL) |>
    readall |>
    s -> split(s,"\r\n") |> # get individual lines
    s -> map(l -> split(l, ","), s)     # decompose each line into
                                        # dates and rates

    maxNumberDates = length(rawSplitted)
    dats = Array(Date{ISOCalendar}, maxNumberDates)
    intRatesStr = Array(String, maxNumberDates)
    skipLines = rep(false, maxNumberDates)

    for ii=1:maxNumberDates
        try
            dats[ii] = date(rawSplitted[ii][1])
            intRatesStr[ii] = rawSplitted[ii][2]
            if intRatesStr[ii] == "ND"
                skipLines[ii] = true
                println("line ", ii, " skipped: no observation")
            end
        catch
            println("line ", ii, " skipped: no date in line")
            skipLines[ii] = true
            continue
        end
    end

    intRatesValid = intRatesStr[!skipLines]
    intRates = Array(Float64, length(intRatesValid))
    for ii=1:length(intRatesValid)
        intRates[ii] = parsefloat(intRatesValid[ii])
    end

    intRatesDf = DataFrame(FEDrate = intRates)
    return Timematr(intRatesDf, dats[!skipLines])
end


#--------------------------------------
## LOAD PACKAGES
using RegNeuralDE, OrdinaryDiffEq, Flux, DiffEqFlux, Tracker
using YAML, Dates, BSON, Random, Statistics, Printf
using CUDA
using RegNeuralDE: accuracy
using Flux: @functor
using Tracker: TrackedReal
import Base.show

CUDA.allowscalar(false)
#--------------------------------------

#--------------------------------------
## CONFIGURATION
config_file = joinpath(pwd(), "experiments", "configs", "mnist_node.yml")
config = YAML.load_file(config_file)

Random.seed!(config["seed"])

hparams = config["hyperparameters"]
BATCH_SIZE = hparams["batch_size"]
EPOCHS = hparams["epochs"]
REGULARIZE = hparams["regularize"]
REG_TYPE = hparams["type"]
STEER = hparams["steer"]
identifier =
    REGULARIZE ? "$(string(now()))_$(REGULARIZE)_$(REG_TYPE)" : "$(string(now()))_vanilla"
EXPERIMENT_LOGDIR = joinpath(pwd(), "results", "mnist_node", identifier)
MODEL_WEIGHTS = joinpath(EXPERIMENT_LOGDIR, "weights.bson")
FILENAME = joinpath(EXPERIMENT_LOGDIR, "results.yml")

# Create a directory to store the results
isdir(EXPERIMENT_LOGDIR) || mkpath(EXPERIMENT_LOGDIR)
cp(config_file, joinpath(EXPERIMENT_LOGDIR, "config.yml"))
#--------------------------------------

#--------------------------------------
## NEURAL NETWORK
# The dynamics of the Neural ODE are time dependent
struct MLPDynamics{D1,D2}
    dense_1::D1
    dense_2::D2
end

@functor MLPDynamics

MLPDynamics(in::Integer, hidden::Integer) =
    MLPDynamics(Dense(in + 1, hidden, CUDA.tanh), Dense(hidden + 1, in, CUDA.tanh))

function (mlp::MLPDynamics)(x::AbstractMatrix, t::TrackedReal)
    _t = CUDA.ones(Float32, 1, size(x, 2)) .* t
    return mlp.dense_2(vcat(mlp.dense_1(vcat(x, _t)), _t))
end
#--------------------------------------

#--------------------------------------
## SETUP THE MODELS + DATASET + TRAINING UTILS
# Get the dataset
train_dataloader, test_dataloader = load_mnist(BATCH_SIZE, x -> cpu(x))

if REG_TYPE == "error_est"
    # Anneal the regularization so that it doesn't overpower the
    # the main objective
    λ₀ = 1.0f2
    λ₁ = 1.0f1
    save_func(u, t, integrator) = integrator.EEst * integrator.dt
    solver = Tsit5()
    global agg = mean
elseif REG_TYPE == "stiff_est"
    # No annealing is generally needed for stiff_est
    λ₀ = 0.1f0
    λ₁ = 0.1f0
    const stability_size =
        Tracker.TrackedReal(1 / Float32(OrdinaryDiffEq.alg_stability_size(Tsit5())))
    function save_func(u, t, integrator)
        stiff_est = abs(integrator.eigen_est)
        return stability_size * ((iszero(stiff_est) || isnan(stiff_est)) ? 0 : stiff_est)
    end
    global agg = maximum
    solver = AutoTsit5(Tsit5())
elseif REG_TYPE == "error_stiff_est"
    λ₀ = 1.0f1
    λ₁ = 1.0f1
    const mul_val = Tracker.TrackedReal(1.0f0)
    const stability_size =
        Tracker.TrackedReal(1 / Float32(OrdinaryDiffEq.alg_stability_size(Tsit5())))
    function save_func(u, t, integrator)
        err_est = integrator.EEst * integrator.dt
        eest = Tracker.data(err_est)
        stiff_est = integrator.eigen_est
        sest = Tracker.data(stiff_est)
        return (
            ((iszero(eest) || isnan(eest)) ? 0 : err_est) +
            0.1f0 * stability_size * ((iszero(sest) || isnan(sest)) ? 0 : stiff_est)
        ) * mul_val
    end
    global agg = mean
    solver = AutoTsit5(Tsit5())
else
    global agg = mean
    solver = Tsit5()
end
b = 0.5f0
sample_tspan_ubound() = [0.0f0, 1.0f0 - (2 * rand(Float32) - 1.0f0) * b]
k = log(λ₀ / λ₁) / EPOCHS
# Exponential Decay
λ_func(t) = λ₀ * exp(-k * t)

# Define the models
mlp_dynamics = MLPDynamics(784, 100)

node = ClassifierNODE(
    Chain(x -> reshape(x, 784, :)) |> track |> gpu,
    TrackedNeuralODE(
        mlp_dynamics |> track |> gpu,
        [0.0f0, 1.0f0],
        true,
        REGULARIZE,
        solver,
        save_everystep = false,
        reltol = 1.4f-8,
        abstol = 1.4f-8,
        save_start = false,
    ),
    Dense(784, 10) |> track |> gpu,
)
ps = Flux.trainable(node)

opt = Flux.Optimise.Optimiser(InvDecay(1.0e-5), Momentum(0.1, 0.9))

function loss_function(x, y, model, p1, p2, p3; λ = 1.0f2, notrack = false)
    tspan = STEER ? sample_tspan_ubound() : nothing
    pred, _, sv = model(x, p1, p2, p3; func = save_func, tspan = tspan)
    cross_entropy = Flux.Losses.logitcrossentropy(pred, y)
    reg = REGULARIZE ? λ * agg(sv.saveval) : zero(eltype(pred))
    total_loss = cross_entropy + reg
    if !notrack
        ce_un = cross_entropy |> untrack
        reg_un = reg |> untrack
        total_loss_un = total_loss |> untrack
        logger(
            false,
            Dict(
                "Total Loss" => total_loss_un,
                "Cross Entropy Loss" => ce_un,
                "Regularization" => reg_un,
            ),
        )
    end
    return total_loss
end
#--------------------------------------

#--------------------------------------
## LOGGING UTILITIES
nfe_counts = Vector{Float64}(undef, EPOCHS + 1)
train_accuracies = Vector{Float64}(undef, EPOCHS + 1)
test_accuracies = Vector{Float64}(undef, EPOCHS + 1)
train_runtimes = Vector{Float64}(undef, EPOCHS + 1)
inference_runtimes = Vector{Float64}(undef, EPOCHS + 1)
train_runtimes[1] = 0

logger = table_logger(
    [
        "Epoch Number",
        "NFE Count",
        "Train Accuracy",
        "Test Accuracy",
        "Train Runtime",
        "Inference Runtime",
    ],
    ["Total Loss", "Cross Entropy Loss", "Regularization"],
)
#--------------------------------------

#--------------------------------------
## RECORD DETAILS BEFORE TRAINING STARTS
dummy_data = train_dataloader.data[1][:, :, :, 1:BATCH_SIZE] |> gpu |> track
stime = time()
_, _nfe, _ = node(dummy_data; func = save_func)
inference_runtimes[1] = time() - stime
train_runtimes[1] = 0.0
nfe_counts[1] = _nfe
train_accuracies[1] = accuracy(node, train_dataloader)
test_accuracies[1] = accuracy(node, test_dataloader)

logger(
    false,
    Dict(),
    0.0,
    nfe_counts[1],
    train_accuracies[1],
    test_accuracies[1],
    train_runtimes[1],
    inference_runtimes[1],
)
#--------------------------------------

#--------------------------------------
## WARMSTART THE GRADIENT COMPUTATION
y_ = zeros(Float32, 10, 1)
y_[1, 1] = 1.0
_ = Tracker.gradient(
    (p1, p2, p3) -> loss_function(
        rand(Float32, 28, 28, 1, 1) |> gpu |> track,
        y_ |> gpu |> track,
        node,
        p1,
        p2,
        p3;
        notrack = true,
    ),
    ps...,
)
#--------------------------------------

#--------------------------------------
## TRAINING
for epoch = 1:EPOCHS
    λ = λ_func(epoch - 1)
    timing = 0

    for (i, (x_, y_)) in enumerate(train_dataloader)
        x = x_ |> gpu |> track
        y = Float32.(y_) |> gpu |> track

        start_time = time()
        gs = Tracker.gradient(
            (p1, p2, p3) -> loss_function(x, y, node, p1, p2, p3; λ = λ),
            ps...,
        )
        update_parameters!(ps, gs, opt)
        timing += time() - start_time

        x = y = nothing
        GC.gc(true)
    end

    # Record the time per epoch
    train_runtimes[epoch+1] = timing

    # Record the NFE count
    start_time = time()
    _, nfe, _ = node(dummy_data)
    inference_runtimes[epoch+1] = time() - start_time
    nfe_counts[epoch+1] = nfe

    # Test and Train Accuracy
    train_accuracies[epoch+1] = accuracy(node, train_dataloader)
    test_accuracies[epoch+1] = accuracy(node, test_dataloader)

    logger(
        false,
        Dict(),
        epoch,
        nfe_counts[epoch+1],
        train_accuracies[epoch+1],
        test_accuracies[epoch+1],
        train_runtimes[epoch+1],
        inference_runtimes[epoch+1],
    )
end
logger(true, Dict())
#--------------------------------------

#--------------------------------------
## STORE THE RESULTS
results = Dict(
    :nfe_counts => nfe_counts,
    :train_accuracies => train_accuracies,
    :test_accuracies => test_accuracies,
    :train_runtimes => train_runtimes,
    :inference_runtimes => inference_runtimes,
)

weights = Flux.params(node) .|> cpu .|> untrack
BSON.@save MODEL_WEIGHTS weights

YAML.write_file(FILENAME, results)
#--------------------------------------


import Base: length

# TODO
# * Implement Hessian and DiagHess

# Define data for logistic regression problem
struct DualLogitData{T <: Real} <: AbstractDataset{T}
    X::AbstractArray{T}
    y::Vector{T}
    x_pred::Vector{T}
end
function DualLogitData(X::Array{T, 2}, y::Vector{T}) where T
    n_feats = size(X, 2)
    return DualLogitData(X, y, zeros(T, n_feats))
end
function DualLogitData(libsmv_file::String; scale_data=false)
    dataset = parse_libsvm(libsmv_file)
    # Convert to dense matrix
    X = to_dense(dataset)
    if scale_data
        scale!(NormalScaler(), X)
    end
    y = dataset.labels
    # Special preprocessing for covtype
    format_label!(y)
    return DualLogitData(X, y)
end

ndata(dat::DualLogitData) = length(dat.y)
nfeatures(dat::DualLogitData) = size(dat.X, 2)

lowerbound(data::DualLogitData) = -0.5 * data.y .- 0.5
upperbound(data::DualLogitData) = -0.5 * data.y .+ 0.5

function predict!(data::DualLogitData, x)
    mul!(data.x_pred, data.X', x)
end

"""
Compute logistic loss of given vector parameter `λ ∈  R^n`.

## Complexity
O(n)

"""
function loss(ω::AbstractVector{T}, data::DualLogitData{T}) where T
    res = zero(T)
    @inbounds for i in 1:ndata(data)
        res += loss(LogLoss(), ω[i], data.y[i])
    end
    return -res
end

"""
Compute gradient of logistic loss for given parameter `λ ∈  R^n`.

## Complexity
O(n)

"""
function gradient!(grad::AbstractVector{T}, ω::AbstractVector{T}, data::DualLogitData{T}) where T
    n = ndata(data)
    invn = one(T) / n
    @inbounds for j in 1:n
        λtmp = ω[j] * data.y[j]
        grad[j] = data.y[j] * (log1p(λtmp) - log(-λtmp))
    end
    return nothing
end

function hessvec!(hv::AbstractVector{T}, ω::AbstractVector{T},
                  vec::AbstractVector{T}, data::DualLogitData{T}) where T
    n = ndata(data)
    invn = one(T) / n
    @inbounds for j in 1:n
        λtmp = ω[j] * data.y[j]
        hv[j] = (-1.0 / λtmp + 1.0 / (1.0 + λtmp)) * vec[j]
    end
    return nothing
end


export iterative

# Efficient least-square solution for affine diagonal component
# using least-square formulation
# function iterative(C::SparseRadialMapComponent, ens::EnsembleState{Nx, Ne}, λ, δ) where {Nx, Ne}
#     @get C (Nx,p)
#     # Compute weights
#     ψ_off, ψ_mono, _ = weights(C, ens)
#
#     # no = size(ψ_off,1)
#     # nd = size(ψ_mono,1)
#     nx = size(ψ_off,1)+size(ψ_mono,1)+1
#     # nlog = size(dψ_mono,1)
#
#     # @assert nd==nlog "Error size of the diag and ∂k weights"
#     # @assert nx==no+nlog+1 "Error size of the weights"
#
#     # Cache for the solution
#     x = zeros(nx)
#
#     # Normalize monotone basis functions
#     μψ = mean(ψ_mono, dims=2)[1,1]
#     σψ = std(ψ_mono, dims=2, corrected=false)[1,1]
#     ψ_mono .-= μψ
#     ψ_mono ./= σψ
#
#     # Normalize off-diagonal covariates
#     σψ_off = std(ψ_off, dims = 2, corrected = false)[:,1]
#     ψ_off ./= σψ_off
#     # Solve for the off-diagonal coefficients and the constant
#     # We can use dψ_mono since it contains only one
#     ψ_at = hcat(ψ_off', ones(Ne))
#
#     # Solve least-square problem
#     # x[1:nx-1] = ψ_a' \ view(ψ_mono,1,:)
#     x0 = view(x,1:nx-1)
#     lsmr!(x0, ψ_at, view(ψ_mono,1,:), λ = λ, atol = 1e-9, btol = 1e-9)
#     # lsmr!(x0, S, view(ψ_mono,1,:), λ = λ)
#
#     @inbounds for i=1:nx-2
#     x[i] *= (-σψ/σψ_off[i])
#     end
#
#     x[nx-1] *= -σψ
#     x[nx-1] -= μψ
#
#     x[end] = √(Ne)/norm(ψ_at*(x[1:nx-1] .* vcat(σψ_off, 1.0)) .+ μψ + σψ*ψ_mono[1,:])
#     x[1:nx-1] .*= x[end]
#     return x
# end

function iterative(C::SparseRadialMapComponent, X, λ, δ)
    NxX, Ne = size(X)
    @get C (Nx,p)

    @assert NxX  == Nx "Wrong dimension of the ensemble matrix `X`"
    # Compute weights
    ψ_off, ψ_mono, _ = compute_weights(C, X)

    # no = size(ψ_off,1)
    # nd = size(ψ_mono,1)
    nx = size(ψ_off,2)+size(ψ_mono,2)+1
    # nlog = size(dψ_mono,1)

    # @assert nd==nlog "Error size of the diag and ∂k weights"
    # @assert nx==no+nlog+1 "Error size of the weights"

    # Cache for the solution
    x = zeros(nx)

    # Normalize monotone basis functions
    μψ = mean(ψ_mono, dims=1)[1,1]
    σψ = std(ψ_mono, dims=1, corrected=false)[1,1]
    ψ_mono .-= μψ
    ψ_mono ./= σψ

    # Normalize off-diagonal covariates
    σψ_off = std(ψ_off, dims = 1, corrected = false)[1,:]
    ψ_off ./= σψ_off'
    # Solve for the off-diagonal coefficients and the constant
    # We can use dψ_mono since it contains only one
    ψ_at = hcat(ψ_off, ones(Ne))

    # Solve least-square problem
    # x[1:nx-1] = ψ_a' \ view(ψ_mono,1,:)
    x0 = view(x,1:nx-1)
    lsmr!(x0, ψ_at, view(ψ_mono,:,1), λ = λ, atol = 1e-9, btol = 1e-9)
    # lsmr!(x0, S, view(ψ_mono,1,:), λ = λ)

    @inbounds for i=1:nx-2
    x[i] *= (-σψ/σψ_off[i])
    end

    x[nx-1] *= -σψ
    x[nx-1] -= μψ

    x[end] = √(Ne)/norm(ψ_at*(x[1:nx-1] .* vcat(σψ_off, 1.0)) .+ μψ + σψ*ψ_mono[:,1])
    x[1:nx-1] .*= x[end]
    return x
end


module TensorKitManifolds

export base, checkbase, isisometry, isunitary
export projecthermitian, projecthermitian!, projectantihermitian, projectantihermitian!
export projectcomplement, projectcomplement!, projectisometric, projectisometric!
export Grassmann, Stiefel, Unitary
export inner, retract, transport, transport!

using TensorKit

# Every submodule -- Grassmann, Stiefel, and Unitary -- implements their own methods for
# these. The signatures should be
# inner(W, Δ₁::Tangent, Δ₂::Tangent; metric)
# retract(W, Δ::Tangent, α::Real; alg)
# transport(Θ::Tangent, W, Δ::Tangent, α::Real, W′; alg)
# where the keyword arguments `alg` and `metric` should always be accepted, even if there is
# only one option for them and they are ignored. The `Tangent` is just a placeholder for the
# tangent type of each manifold. Similarly each submodule defines a `project!` function,
# which too should accept a keyword argument `metric`, even if it is ignored.
function inner end
function retract end
function transport end
function transport! end
function base end
function checkbase end
checkbase(x, y, z, args...) = checkbase(checkbase(x, y), z, args...)

# the machine epsilon for the elements of an object X, name inspired from eltype
eleps(X) = eps(real(eltype(X)))

function isisometry(W::AbstractTensorMap; tol = 10*eleps(W))
    WdW = W'*W
    s = zero(float(real(eltype(W))))
    for (c,b) in blocks(WdW)
        _subtractone!(b)
        s += dim(c)*length(b)
    end
    return norm(WdW) <= tol*sqrt(s)
end

isunitary(W::AbstractTensorMap; tol = 10*eleps(W)) =
    isisometry(W; tol = tol) && isisometry(W'; tol = tol)

function projecthermitian!(W::AbstractTensorMap)
    codomain(W) == domain(W) ||
        throw(DomainError("Tensor with distinct domain and codomain cannot be hermitian."))
    for (c,b) in blocks(W)
        _projecthermitian!(b)
    end
    return W
end
function projectantihermitian!(W::AbstractTensorMap)
    codomain(W) == domain(W) ||
        throw(DomainError("Tensor with distinct domain and codomain cannot be anithermitian."))
    for (c,b) in blocks(W)
        _projectantihermitian!(b)
    end
    return W
end

struct PolarNewton <: TensorKit.OrthogonalFactorizationAlgorithm
end
function projectisometric!(W::AbstractTensorMap; alg = Polar())
    if alg isa TensorKit.Polar || alg isa TensorKit.SDD
        foreach(blocks(W)) do (c,b)
            _polarsdd!(b)
        end
    elseif alg isa TensorKit.SVD
        foreach(blocks(W)) do (c,b)
            _polarsvd!(b)
        end
    elseif alg isa PolarNewton
        foreach(blocks(W)) do (c,b)
            _polarnewton!(b)
        end
    else
        throw(ArgumentError("unkown algorithm for projectisometric!: alg = $alg"))
    end
    return W
end

function projectcomplement!(X::AbstractTensorMap, W::AbstractTensorMap;
                            tol = 10*eleps(X))
    P = W'*X
    nP = norm(P)
    nX = norm(X)
    dP = dim(P)
    while nP > tol*max(dP, nX)
        X = mul!(X, W, P, -1, 1)
        P = W'*X
        nP = norm(P)
    end
    return X
end

projecthermitian(W::AbstractTensorMap) = projecthermitian!(copy(W))
projectantihermitian(W::AbstractTensorMap) = projectantihermitian!(copy(W))

function projectisometric(W::AbstractTensorMap;
                          alg::TensorKit.OrthogonalFactorizationAlgorithm = Polar())
    return projectisometric!(copy(W); alg=alg)
end
function projectcomplement(X::AbstractTensorMap, W::AbstractTensorMap,
                           tol = 10*eleps(X))
    return projectcomplement!(copy(X), W; tol=tol)
end

include("auxiliary.jl")
include("grassmann.jl")
include("stiefel.jl")
include("unitary.jl")

end


using LinearAlgebra, SparseArrays, Random, QPDAS
# Number of variables
n = 1000
# Number of halfspaces
mi = 50 # Inequalities
me = 50 # Equalities

Random.seed!(1)
# One point in polytope
x0 = randn(n)

# Create polytope containing x0
# Inequality
C = Matrix{Float64}(undef, mi, n)
d = randn(mi)

# Make sure x0 is in polytope by setting sign of inequality
for i = 1:mi
    v = randn(n)
    b = randn()
    if v'x0  <= b
        C[i,:] .= v
    else
        C[i,:] .= -v
    end
    d[i] = b
end

# Create equality
A = randn(me, n)
b = A*x0

###### Solve with only inequality

# Project from here
x = randn(n)

@time qp = QuadraticProgram(zeros(0,n), zeros(0), C, d, -x, I)
@time sol, val = solve!(qp)
# Test Feasibility
maximum(C*sol-d)

###### Solve with equality and inequality

# Project from here
x = randn(n)
qp = QuadraticProgram(A, b, C, d, -x, I)
solve!(qp)
# Test Feasibility
maximum(C*sol-d)


"""
 M6 function
 Quintic spline
               |  (3-q)^5-6(2-q)^5+15(1-q)^5  for 0 <= q < 1
 M6(x) = 1/120 |  (3-q)^5-6(2-q)^5            for 1 <= q < 2
               |  (3-q)^5                     for 2 <= q < 3
               |  0                           for q >= 3
"""

@inline function f_m6( q :: Float64 )

    if ( q < 1.0 )
        f_m6 = (3-q)^5-6*(2-q)^5+15*(1-q)^5
    elseif ( q >= 1.0 && q < 2.0 )
        f_m6 = (3-q)^5-6*(2-q)^5
    elseif ( q >= 2 && q < 3 )
        f_m6 = (3-q)^5
    else
	f_m6 = 0.0
    end

    f_m6 / 120

end

export compute_rho_m6!

function compute_rho_m6!( fields    :: MeshFields,
                          particles :: Particles,
                          xt        :: Array{ComplexF64,3},
                          ua        :: UA )

    fill!(fields.ρ , 0.0)
    nx, ny = fields.mesh.nx, fields.mesh.ny
    dx, dy = fields.mesh.dx, fields.mesh.dy

    xmin, xmax = fields.mesh.xmin, fields.mesh.xmax
    ymin, ymax = fields.mesh.ymin, fields.mesh.ymax
    dimx = xmax - xmin
    dimy = ymax - ymin

    for k = 1:particles.nbpart

        t = particles.t[k]

        mul!(ua.ftau, ua.ptau, view(xt,:,1,k))

        for n = 1:ua.ntau
            ua.ftau[n] *= exp(1im*ua.ltau[n]*t/ua.ε)/ua.ntau
        end

        xt1 = real(sum(ua.ftau))

        mul!(ua.ftau, ua.ptau, view(xt,:,2,k))

        for n = 1:ua.ntau
            ua.ftau[n] *= exp.(1im*ua.ltau[n]*t/ua.ε)/ua.ntau
        end

        xt2 = real(sum(ua.ftau))

        xt1 = mod( xt1 - xmin , dimx)
        xt2 = mod( xt2 - ymin , dimy)

        px = xt1/dx
        py = xt2/dy

        particles.x[1,k] = xt1 + xmin
        particles.x[2,k] = xt2 + ymin

        i  = trunc(Int32, px)
        j  = trunc(Int32, py)
        dpx = px - i
        dpy = py - j

        weight = particles.w

        im3 = mod(i-3,nx)+1
        im2 = mod(i-2,nx)+1
        im1 = mod(i-1,nx)+1
        ip1 = mod(i+1,nx)+1
        ip2 = mod(i+2,nx)+1
        ip3 = mod(i+3,nx)+1
        jm3 = mod(j-3,ny)+1
        jm2 = mod(j-2,ny)+1
        jm1 = mod(j-1,ny)+1
        jp1 = mod(j+1,ny)+1
        jp2 = mod(j+2,ny)+1
        jp3 = mod(j+3,ny)+1

        i = i+1
        j = j+1

        cm3x = f_m6(3+dpx)
        cp3x = f_m6(3-dpx)
        cm2x = f_m6(2+dpx)
        cp2x = f_m6(2-dpx)
        cm1x = f_m6(1+dpx)
        cp1x = f_m6(1-dpx)
        cx   = f_m6(dpx)
        cy   = f_m6(dpy)
        cp1y = f_m6(1-dpy)
        cm1y = f_m6(1+dpy)
        cp2y = f_m6(2-dpy)
        cm2y = f_m6(2+dpy)
        cp3y = f_m6(3-dpy)
        cm3y = f_m6(3+dpy)

	    fields.ρ[im3,jm3] += cm3x * cm3y * weight
        fields.ρ[im3,jm2] += cm3x * cm2y * weight
        fields.ρ[im3,jm1] += cm3x * cm1y * weight
        fields.ρ[im3,j  ] += cm3x * cy   * weight
        fields.ρ[im3,jp1] += cm3x * cp1y * weight
        fields.ρ[im3,jp2] += cm3x * cp2y * weight
        fields.ρ[im3,jp3] += cm3x * cp3y * weight

	    fields.ρ[im2,jm3] += cm2x * cm3y * weight
        fields.ρ[im2,jm2] += cm2x * cm2y * weight
        fields.ρ[im2,jm1] += cm2x * cm1y * weight
        fields.ρ[im2,j  ] += cm2x * cy   * weight
        fields.ρ[im2,jp1] += cm2x * cp1y * weight
        fields.ρ[im2,jp2] += cm2x * cp2y * weight
        fields.ρ[im2,jp3] += cm2x * cp3y * weight

	    fields.ρ[im1,jm3] += cm1x * cm3y * weight
        fields.ρ[im1,jm2] += cm1x * cm2y * weight
        fields.ρ[im1,jm1] += cm1x * cm1y * weight
        fields.ρ[im1,j  ] += cm1x * cy   * weight
        fields.ρ[im1,jp1] += cm1x * cp1y * weight
        fields.ρ[im1,jp2] += cm1x * cp2y * weight
        fields.ρ[im1,jp3] += cm1x * cp3y * weight

        fields.ρ[i  ,jm3] += cx   * cm3y * weight
        fields.ρ[i  ,jm2] += cx   * cm2y * weight
        fields.ρ[i  ,jm1] += cx   * cm1y * weight
        fields.ρ[i  ,j  ] += cx   * cy   * weight
        fields.ρ[i  ,jp1] += cx   * cp1y * weight
        fields.ρ[i  ,jp2] += cx   * cp2y * weight
        fields.ρ[i  ,jp3] += cx   * cp3y * weight

        fields.ρ[ip1,jm3] += cp1x * cm3y * weight
        fields.ρ[ip1,jm2] += cp1x * cm2y * weight
        fields.ρ[ip1,jm1] += cp1x * cm1y * weight
        fields.ρ[ip1,j  ] += cp1x * cy   * weight
        fields.ρ[ip1,jp1] += cp1x * cp1y * weight
        fields.ρ[ip1,jp2] += cp1x * cp2y * weight
        fields.ρ[ip1,jp3] += cp1x * cp3y * weight

        fields.ρ[ip2,jm3] += cp2x * cm3y * weight
        fields.ρ[ip2,jm2] += cp2x * cm2y * weight
        fields.ρ[ip2,jm1] += cp2x * cm1y * weight
        fields.ρ[ip2,j  ] += cp2x * cy   * weight
        fields.ρ[ip2,jp1] += cp2x * cp1y * weight
        fields.ρ[ip2,jp2] += cp2x * cp2y * weight
        fields.ρ[ip2,jp3] += cp2x * cp3y * weight

        fields.ρ[ip3,jm3] += cp3x * cm3y * weight
        fields.ρ[ip3,jm2] += cp3x * cm2y * weight
        fields.ρ[ip3,jm1] += cp3x * cm1y * weight
        fields.ρ[ip3,j  ] += cp3x * cy   * weight
        fields.ρ[ip3,jp1] += cp3x * cp1y * weight
        fields.ρ[ip3,jp2] += cp3x * cp2y * weight
        fields.ρ[ip3,jp3] += cp3x * cp3y * weight

    end

    fields.ρ[1:nx,ny+1] .= fields.ρ[1:nx,1]
    fields.ρ[nx+1,1:ny] .= fields.ρ[1,1:ny]
    fields.ρ[nx+1,ny+1]  = fields.ρ[1,1]

    fields.ρ ./= (dx*dy)

    rho_total = sum(view(fields.ρ,1:nx,1:ny)) * dx * dy

    fields.ρ .-= rho_total/dimx/dimy

end

function compute_rho_m6!( fields  :: MeshFields, particles :: Particles)

    fill!(fields.ρ , 0.0)
    nx = fields.mesh.nx
    ny = fields.mesh.ny

    dx = fields.mesh.dx
    dy = fields.mesh.dy

    xmin, xmax = fields.mesh.xmax, fields.mesh.xmin
    ymin, ymax = fields.mesh.ymax, fields.mesh.ymin
    dimx = fields.mesh.xmax - fields.mesh.xmin
    dimy = fields.mesh.ymax - fields.mesh.ymin

    for k = 1:particles.nbpart

        xt1 = mod( particles.x[1,k] - xmin , dimx)
        xt2 = mod( particles.x[2,k] - ymin , dimy)

        particles.x[1,k] = xt1 + xmin
        particles.x[2,k] = xt2 + ymin

        px = xt1/dx
        py = xt2/dy

        i = trunc(Int32, px)
        j = trunc(Int32, py)

        dpx = px - i
        dpy = py - j

        weight = particles.w

        im3 = mod(i-3,nx)+1
        im2 = mod(i-2,nx)+1
        im1 = mod(i-1,nx)+1
        ip1 = mod(i+1,nx)+1
        ip2 = mod(i+2,nx)+1
        ip3 = mod(i+3,nx)+1
        jm3 = mod(j-3,ny)+1
        jm2 = mod(j-2,ny)+1
        jm1 = mod(j-1,ny)+1
        jp1 = mod(j+1,ny)+1
        jp2 = mod(j+2,ny)+1
        jp3 = mod(j+3,ny)+1

        i = i+1
        j = j+1

        cm3x = f_m6(3+dpx)
        cp3x = f_m6(3-dpx)
        cm2x = f_m6(2+dpx)
        cp2x = f_m6(2-dpx)
        cm1x = f_m6(1+dpx)
        cp1x = f_m6(1-dpx)
        cx   = f_m6(dpx)
        cy   = f_m6(dpy)
        cp1y = f_m6(1-dpy)
        cm1y = f_m6(1+dpy)
        cp2y = f_m6(2-dpy)
        cm2y = f_m6(2+dpy)
        cp3y = f_m6(3-dpy)
        cm3y = f_m6(3+dpy)

	    fields.ρ[im3,jm3] += cm3x * cm3y * weight
        fields.ρ[im3,jm2] += cm3x * cm2y * weight
        fields.ρ[im3,jm1] += cm3x * cm1y * weight
        fields.ρ[im3,j  ] += cm3x * cy   * weight
        fields.ρ[im3,jp1] += cm3x * cp1y * weight
        fields.ρ[im3,jp2] += cm3x * cp2y * weight
        fields.ρ[im3,jp3] += cm3x * cp3y * weight

	    fields.ρ[im2,jm3] += cm2x * cm3y * weight
        fields.ρ[im2,jm2] += cm2x * cm2y * weight
        fields.ρ[im2,jm1] += cm2x * cm1y * weight
        fields.ρ[im2,j  ] += cm2x * cy   * weight
        fields.ρ[im2,jp1] += cm2x * cp1y * weight
        fields.ρ[im2,jp2] += cm2x * cp2y * weight
        fields.ρ[im2,jp3] += cm2x * cp3y * weight

	    fields.ρ[im1,jm3] += cm1x * cm3y * weight
        fields.ρ[im1,jm2] += cm1x * cm2y * weight
        fields.ρ[im1,jm1] += cm1x * cm1y * weight
        fields.ρ[im1,j  ] += cm1x * cy   * weight
        fields.ρ[im1,jp1] += cm1x * cp1y * weight
        fields.ρ[im1,jp2] += cm1x * cp2y * weight
        fields.ρ[im1,jp3] += cm1x * cp3y * weight

        fields.ρ[i  ,jm3] += cx   * cm3y * weight
        fields.ρ[i  ,jm2] += cx   * cm2y * weight
        fields.ρ[i  ,jm1] += cx   * cm1y * weight
        fields.ρ[i  ,j  ] += cx   * cy   * weight
        fields.ρ[i  ,jp1] += cx   * cp1y * weight
        fields.ρ[i  ,jp2] += cx   * cp2y * weight
        fields.ρ[i  ,jp3] += cx   * cp3y * weight

        fields.ρ[ip1,jm3] += cp1x * cm3y * weight
        fields.ρ[ip1,jm2] += cp1x * cm2y * weight
        fields.ρ[ip1,jm1] += cp1x * cm1y * weight
        fields.ρ[ip1,j  ] += cp1x * cy   * weight
        fields.ρ[ip1,jp1] += cp1x * cp1y * weight
        fields.ρ[ip1,jp2] += cp1x * cp2y * weight
        fields.ρ[ip1,jp3] += cp1x * cp3y * weight

        fields.ρ[ip2,jm3] += cp2x * cm3y * weight
        fields.ρ[ip2,jm2] += cp2x * cm2y * weight
        fields.ρ[ip2,jm1] += cp2x * cm1y * weight
        fields.ρ[ip2,j  ] += cp2x * cy   * weight
        fields.ρ[ip2,jp1] += cp2x * cp1y * weight
        fields.ρ[ip2,jp2] += cp2x * cp2y * weight
        fields.ρ[ip2,jp3] += cp2x * cp3y * weight

        fields.ρ[ip3,jm3] += cp3x * cm3y * weight
        fields.ρ[ip3,jm2] += cp3x * cm2y * weight
        fields.ρ[ip3,jm1] += cp3x * cm1y * weight
        fields.ρ[ip3,j  ] += cp3x * cy   * weight
        fields.ρ[ip3,jp1] += cp3x * cp1y * weight
        fields.ρ[ip3,jp2] += cp3x * cp2y * weight
        fields.ρ[ip3,jp3] += cp3x * cp3y * weight

    end

    fields.ρ[1:nx,ny+1] .= fields.ρ[1:nx,1]
    fields.ρ[nx+1,1:ny] .= fields.ρ[1,1:ny]
    fields.ρ[nx+1,ny+1]  = fields.ρ[1,1]

    fields.ρ ./= (dx*dy)

    rho_total = sum(fields.ρ[1:nx,1:ny]) * dx * dy

    println( " rho_total = $rho_total ")

    fields.ρ .-= rho_total/dimx/dimy

end


# ## Data

import CSV
import Random
import TuringModels

using DataFrames

Random.seed!(1)

data_path = joinpath(TuringModels.project_root, "data", "UCBadmit.csv")
df = CSV.read(data_path, DataFrame; delim=';')

# ## Model

using Turing

dept_map = Dict(key => idx for (idx, key) in enumerate(unique(df.dept)))
df.male = [g == "male" ? 1 : 0 for g in df.gender]
df.dept_id = [dept_map[de] for de in df.dept]

@model m13_2(applications, dept_id, male, admit) = begin
    sigma_dept ~ truncated(Cauchy(0, 2), 0, Inf)
    bm ~ Normal(0, 1)
    a ~ Normal(0, 10)
    a_dept ~ filldist(Normal(a, sigma_dept), 6)

    logit_p = a_dept[dept_id] + bm*male

    admit .~ BinomialLogit.(applications, logit_p)
end

chns = sample(
    m13_2(df.applications, df.dept_id, df.male, df.admit),
    NUTS(),
    1000
)

# ## Output

# \defaultoutput{}

# ## Original output

"""
Inference for Stan model: 359c2483e3bdbf74fd0484be27c2909b.
    3 chains, each with iter=4500; warmup=500; thin=1;
    post-warmup draws per chain=4000, total post-warmup draws=12000.

                   mean se_mean   sd     2.5%      25%      50%      75%    97.5%
    a_dept[1]      0.67    0.00 0.10     0.48     0.61     0.67     0.74     0.87
    a_dept[2]      0.63    0.00 0.12     0.40     0.55     0.63     0.71     0.85
    a_dept[3]     -0.58    0.00 0.08    -0.73    -0.63    -0.58    -0.53    -0.44
    a_dept[4]     -0.62    0.00 0.09    -0.79    -0.67    -0.62    -0.56    -0.45
    a_dept[5]     -1.06    0.00 0.10    -1.26    -1.13    -1.06    -0.99    -0.86
    a_dept[6]     -2.61    0.00 0.16    -2.92    -2.71    -2.60    -2.50    -2.31
    a             -0.59    0.01 0.66    -1.94    -0.97    -0.59    -0.21     0.73
    bm            -0.09    0.00 0.08    -0.25    -0.15    -0.09    -0.04     0.06
    sigma_dept     1.49    0.01 0.62     0.78     1.09     1.35     1.71     3.01
    lp__       -2602.27    0.04 2.27 -2607.74 -2603.48 -2601.90 -2600.61 -2598.95
               n_eff Rhat
    a_dept[1]   8191    1
    a_dept[2]   8282    1
    a_dept[3]  10974    1
    a_dept[4]  10034    1
    a_dept[5]  10783    1
    a_dept[6]  10344    1
    a           5839    1
    bm          7137    1
    sigma_dept  4626    1
    lp__        4200    1
""";


using LazyInverses
using LinearAlgebra
using BenchmarkTools

n = 1024
x = randn(n)
A = randn(n, n)
@time B = inv(A)
@time C = inverse(A)
println(A \ x ≈ C * x)
println(A * x ≈ C \ x)


#clustered nearest neaighbor
function Kclearest_neighbor{T<:Real}(distmat::Matrix{T},n::Int,
                                    firstcity::Int = 0)
  scan = kmeans(distmat,n)
  asgn = scan.assignments
  bitClust = trues(n)
	numCities = size(distmat, 1)
  bitNode = trues(numCities)
  master = collect(1:numCities)
	# put first city on path
	path = Vector{Int}()
  sizehint!(path,numCities)
	if firstcity == 0
		firstcity = rand(1:numCities)
	end
	push!(path, firstcity)
	# cities to visit
  while in(true,bitNode)
    curCity = path[end]
    bitNode[curCity] = false
    Cid = asgn[curCity]
    if Cid == 0
      citiesToVisit = master[bitNode]
    elseif bitClust[Cid]
      bitPals = asgn.==Cid
      citiesToVisit = master[bitNode&bitPals]
      bitClust[Cid] = false
    else
      citiesToVisit = master[bitNode]
    end
	# nearest neighbor loop
  	while asgn[curCity] == Cid && !isempty(citiesToVisit)
      dists = distmat[citiesToVisit,curCity]
  		_, nextInd = findmin(dists)
  		nextCity = citiesToVisit[nextInd]
  		push!(path, nextCity)
      bitNode[nextCity] = false
  		deleteat!(citiesToVisit, nextInd)
      curCity = path[end]
  	end
  end

	return path
end


struct Cost{T}
    #TODO: types for methods
    val
    grad
    hess
    sp::Vector{Vector{Int}}
    val_cache::Vector{T}
    grad_cache::Vector{T}
    hess_cache::Vector{T}
    idx::Vector{Int}
end

function Cost(f::Function, nx::Int, nu::Int, nw::Int, idx::Vector{Int}; eval_hess=false)
    #TODO: option to load/save methods
    @variables x[1:nx], u[1:nu], w[1:nw]
    val = f(x, u, w)
    grad = Symbolics.gradient(val, [x; u])
    val_func = eval(Symbolics.build_function([val], x, u, w)[2])
    grad_func = eval(Symbolics.build_function(grad, x, u, w)[2])
    if eval_hess
        hess = Symbolics.sparsehessian(val, [x; u])
        hess_func = eval(Symbolics.build_function(hess.nzval, x, u, w)[2])
        sparsity = [findnz(hess)[1:2]...]
        nh = length(hess.nzval)
    else
        hess_func = Expr(:null)
        sparsity = [Int[]]
        nh = 0
    end

    return Cost(val_func, grad_func, hess_func, sparsity,
        zeros(1), zeros(nx + nu), zeros(nh), idx)
end

Objective{T} = Vector{Cost{T}} where T

function eval_obj(obj::Objective, x, u, w)
    J = 0.0
    for cost in obj
        for t in cost.idx
            cost.val(cost.val_cache, x[t], u[t], w[t])
            J += cost.val_cache[1]
        end
    end
    return J
end

function eval_obj_grad!(grad, idx, obj::Objective, x, u, w)
    for cost in obj
        for t in cost.idx
            cost.grad(cost.grad_cache, x[t], u[t], w[t])
            @views grad[idx[t]] .+= cost.grad_cache
            # fill!(cost.grad_cache, 0.0) # TODO: confirm this is necessary
        end
    end
end

#TODO: test
function eval_obj_hess!(hess, idx, obj::Objective, x, u, w, σ)
    i = 1
    for cost in obj
        for t in cost.idx
            cost.hess(cost.hess_cache, x[t], u[t], w[t])
            cost.hess_cache .*= σ
            @views hess[idx[i]] .+= cost.hess_cache
            i += 1
            fill!(cost.hess_cache, 0.0) # TODO: confirm this is necessary
        end
    end
end

function sparsity_hessian(obj::Objective, nx::Vector{Int}, nu::Vector{Int})
    row = Int[]
    col = Int[]
    for cost in obj
        if !isempty(cost.sp[1])
            for t in cost.idx
                shift = (t > 1 ? (sum(nx[1:t-1]) + sum(nu[1:t-1])) : 0)
                push!(row, (cost.sp[1] .+ shift)...)
                push!(col, (cost.sp[2] .+ shift)...)
            end
        end
    end
    return collect(zip(row, col))
end

function hessian_indices(obj::Objective, key::Vector{Tuple{Int,Int}}, nx::Vector{Int}, nu::Vector{Int})
    idx = Vector{Int}[]
    for cost in obj
        if !isempty(cost.sp[1])
            for t in cost.idx
                row = Int[]
                col = Int[]
                shift = (t > 1 ? (sum(nx[1:t-1]) + sum(nu[1:t-1])) : 0)
                push!(row, (cost.sp[1] .+ shift)...)
                push!(col, (cost.sp[2] .+ shift)...)
                rc = collect(zip(row, col))
                push!(idx, [findfirst(x -> x == i, key) for i in rc])
            end
        end
    end
    return idx
end


## the learning algorithm: PLS1 - single target
function trainer(model::PLS1Model{T},
                 X::AbstractArray{T}, Y::Vector{T}) where T<:AbstractFloat

    W,b,P  = model.W,model.b,model.P
    nfactors = model.nfactors
    for i = 1:nfactors
        W[:,i] = X'Y
        W[:,i] /= norm(W[:,i])#sqrt.(W[:,i]'*W[:,i])
        R      = X*W[:,i]
        Rn     = R'/(R'R) # change to use function...
        P[:,i] = Rn*X
        b[i]   = Rn * Y
        if abs(b[i]) <= 1e-3
           print("PLS1 converged. No need learning with more than $(i) factors")
           model.nfactors = i
           break
        end
        X      = X - R * P[:,i]'
        Y      = Y - R * b[i]
    end

    return model
end

function predictor(model::PLS1Model{T},
                   X::AbstractArray{T}) where T<:AbstractFloat

    W,b,P    = model.W,model.b,model.P
    nfactors = model.nfactors
    nrows    = size(X,1)
    Y = zeros(T,nrows)

    for i = 1:nfactors
        R      = X*W[:,i]
        Y      = Y + R*b[i]
        X      = X - R*P[:,i]'
    end

    return Y
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    Ball(center, radius)

A Ball with `center` and `radius`.
"""
struct Ball{Dim,T} <: Primitive{Dim,T}
  center::Point{Dim,T}
  radius::T
end

Ball(center::Tuple, radius) = Ball(Point(center), radius)

paramdim(::Type{<:Ball{Dim}}) where {Dim} = Dim

center(b::Ball) = b.center
radius(b::Ball) = b.radius

# https://en.wikipedia.org/wiki/Volume_of_an_n-ball
function measure(b::Ball{Dim}) where {Dim}
  r, n = b.radius, Dim
  (π^(n/2) * r^n) / gamma(n/2 + 1)
end

function Base.in(p::Point, b::Ball)
  x = coordinates(p)
  c = coordinates(b.center)
  r = b.radius
  sum(abs2, x - c) ≤ r^2
end

boundary(b::Ball) = Sphere(b.center, b.radius)


# This definitely needs cleaning
using OrdinaryDiffEq, DiffEqDevTools, DiffEqBase, DiffEqProblemLibrary, Base.Test
probArr = Vector{ODETestProblem}(2)
probArr[1] = prob_ode_linear

probArr[2] = prob_ode_2Dlinear
srand(100)
## Convergence Testing
println("Convergence Test on Linear")
dts = 1.//2.^(8:-1:4)
testTol = 0.2

for i = 1:2
  prob = probArr[i]
  println("Special RKs")
  sim = test_convergence(dts,prob,Euler())
  @test abs(sim.𝒪est[:final]-1) < testTol
  sim2 = test_convergence(dts,prob,Midpoint())
  @test abs(sim2.𝒪est[:l∞]-2) < testTol
  sim3 = test_convergence(dts,prob,RK4())
  @test abs(sim3.𝒪est[:l∞]-4) < testTol
  sim4 = test_convergence(dts,prob,BS3())
  @test abs(sim4.𝒪est[:l2]-3) < testTol

  ### Stiff Solvers

  println("Convergence Test on Stiff")
  dts = 1.//2.^(8:-1:4)

  sim12 = test_convergence(dts,prob,ImplicitEuler(autodiff=true))
  @test abs(sim12.𝒪est[:final]-1) < testTol
  sim122 = test_convergence(dts,prob,ImplicitEuler(autodiff=false))
  @test abs(sim122.𝒪est[:final]-1) < testTol
  sim13 = test_convergence(dts,prob,Trapezoid(autodiff=true))
  @test abs(sim13.𝒪est[:final]-2) < testTol
  sim132 = test_convergence(dts,prob,Trapezoid(autodiff=false))
  @test abs(sim132.𝒪est[:final]-2) < testTol
end


using Plots

include("rw_1d_stat.jl")

num = length(ARGS)>0 ? parse(ARGS[1]) : 1000
nsteps = length(ARGS)>1 ? parse(ARGS[2]) : 1000
prob = length(ARGS)>2 ? parse(ARGS[3]) : 0.5

means, vars = randomwalk(num, nsteps, prob)

pl = plot(means, xaxis=("t"), yaxis=("E[x]"), color=[:black], legend=false)
savefig(pl, "mean.pdf")
pl = plot(vars, xaxis=("t"), yaxis=("V[x]"), color=[:black], legend=false)
savefig(pl, "var.pdf")


inpt = """
Alice would gain 2 happiness units by sitting next to Bob.
Alice would gain 26 happiness units by sitting next to Carol.
Alice would lose 82 happiness units by sitting next to David.
Alice would lose 75 happiness units by sitting next to Eric.
Alice would gain 42 happiness units by sitting next to Frank.
Alice would gain 38 happiness units by sitting next to George.
Alice would gain 39 happiness units by sitting next to Mallory.
Bob would gain 40 happiness units by sitting next to Alice.
Bob would lose 61 happiness units by sitting next to Carol.
Bob would lose 15 happiness units by sitting next to David.
Bob would gain 63 happiness units by sitting next to Eric.
Bob would gain 41 happiness units by sitting next to Frank.
Bob would gain 30 happiness units by sitting next to George.
Bob would gain 87 happiness units by sitting next to Mallory.
Carol would lose 35 happiness units by sitting next to Alice.
Carol would lose 99 happiness units by sitting next to Bob.
Carol would lose 51 happiness units by sitting next to David.
Carol would gain 95 happiness units by sitting next to Eric.
Carol would gain 90 happiness units by sitting next to Frank.
Carol would lose 16 happiness units by sitting next to George.
Carol would gain 94 happiness units by sitting next to Mallory.
David would gain 36 happiness units by sitting next to Alice.
David would lose 18 happiness units by sitting next to Bob.
David would lose 65 happiness units by sitting next to Carol.
David would lose 18 happiness units by sitting next to Eric.
David would lose 22 happiness units by sitting next to Frank.
David would gain 2 happiness units by sitting next to George.
David would gain 42 happiness units by sitting next to Mallory.
Eric would lose 65 happiness units by sitting next to Alice.
Eric would gain 24 happiness units by sitting next to Bob.
Eric would gain 100 happiness units by sitting next to Carol.
Eric would gain 51 happiness units by sitting next to David.
Eric would gain 21 happiness units by sitting next to Frank.
Eric would gain 55 happiness units by sitting next to George.
Eric would lose 44 happiness units by sitting next to Mallory.
Frank would lose 48 happiness units by sitting next to Alice.
Frank would gain 91 happiness units by sitting next to Bob.
Frank would gain 8 happiness units by sitting next to Carol.
Frank would lose 66 happiness units by sitting next to David.
Frank would gain 97 happiness units by sitting next to Eric.
Frank would lose 9 happiness units by sitting next to George.
Frank would lose 92 happiness units by sitting next to Mallory.
George would lose 44 happiness units by sitting next to Alice.
George would lose 25 happiness units by sitting next to Bob.
George would gain 17 happiness units by sitting next to Carol.
George would gain 92 happiness units by sitting next to David.
George would lose 92 happiness units by sitting next to Eric.
George would gain 18 happiness units by sitting next to Frank.
George would gain 97 happiness units by sitting next to Mallory.
Mallory would gain 92 happiness units by sitting next to Alice.
Mallory would lose 96 happiness units by sitting next to Bob.
Mallory would lose 51 happiness units by sitting next to Carol.
Mallory would lose 81 happiness units by sitting next to David.
Mallory would gain 31 happiness units by sitting next to Eric.
Mallory would lose 73 happiness units by sitting next to Frank.
Mallory would lose 89 happiness units by sitting next to George."""

names = Vector{ASCIIString}()
for line in split(inpt, "\n")
	items = split(line[1:end-1]) # get rid of the . at the end
	if !(items[1] in names)
		push!(names, items[1])
	end
	if !(items[end] in names)
		push!(names, items[end])
	end
end
push!(names, "Jan")
happiness = zeros(Int64, length(names), length(names))

for line in split(inpt, "\n")
	items = split(line[1:end-1]) # get rid of the . at the end
	happy = parse(Int, items[4])
	if items[3] == "lose"
		happy *= -1
	end
	i = findfirst(names, items[1])
	j = findfirst(names, items[end])
	happiness[i, j] = happy
end

function tableHappiness(arrangement)
	sum = 0
	# first wrap around
	i = findfirst(names, arrangement[1])
	j = findfirst(names, arrangement[end])
	sum += happiness[i,j] + happiness[j,i]
	for a in 1:length(arrangement)-1
		i = findfirst(names, arrangement[a])
		j = findfirst(names, arrangement[a+1])
		sum += happiness[i,j] + happiness[j,i]
	end
	return sum
end

best = (0, names)
for p in permutations(names)
	h = tableHappiness(p)
	if h > best[1]
		best = (h, p)
	end
end
println(best)


#*************************************************************************#
#
# Preprocessing data
#
#*************************************************************************#
function _compute_stats(mdl::StaticModel,
                        yterms::Vector{Vector{Int}},
                        uterms::Vector{Vector{Int}},
                        iddata::IdData)
    # Get std and mean
    y_std = std(iddata.y, 2)
    u_std = std(iddata.u, 2)
    y_mean = mean(iddata.y, 2)
    u_mean = mean(iddata.u, 2)

    # Compute mean μ and standard deviation σ
    # for each element of the model input x
    ninputs = length_input(mdl)
    k = 1
    σ_x = Vector{Float64}(ninputs)
    μ_x  = Vector{Float64}(ninputs)
    for (i, ynterms) in Enumerate(yterms)
        for yn in ynterms
            σ_x[k] = y_std[i]
            μ_x[k] = y_mean[i]
            k += 1
        end
    end

    for (i, unterms) in Enumerate(uterms)
        for un in unterms
            σ_x[k] = u_std[i]
            μ_x[k] = u_mean[i]
            k += 1
        end
    end

    # Compute mean μ and standard deviation σ
    # for each element of the model input z
    σ_z = vec(y_std)
    μ_z = vec(y_mean)

    return σ_x, μ_x, σ_z, μ_z
end

"""
    learn_offset(mdl, yterms, uterms, iddata) -> new_mdl

Learn offset from data and generate new model with the given offset
embbeded.
"""
function learn_offset(mdl::StaticModel,
                      yterms::Vector{Vector{Int}},
                      uterms::Vector{Vector{Int}},
                      iddata::IdData)

    σ_x, μ_x, σ_z, μ_z = _compute_stats(mdl, yterms, uterms, iddata)

    # Build New Series Model
    new_model = SeriesModel([AffineMap(ones(σ_x), -μ_x), mdl, AffineMap(ones(σ_z), μ_z)])

    return new_model
end

"""
    learn_normalization(mdl, yterms, uterms, iddata[, nσ]) -> new_mdl

Learn offset and scale from data and generate new model with the given offset
and scale embbeded. Where `nσ` is the number of standard deviations the data
will be internally divided by (default is 3).
"""
function learn_normalization{R<:Real}(mdl::StaticModel,
                                      yterms::Vector{Vector{Int}},
                                      uterms::Vector{Vector{Int}},
                                      iddata::IdData;
                                      nσ::R=3.0)

    σ_x, μ_x, σ_z, μ_z = _compute_stats(mdl, yterms, uterms, iddata)

    # Replace σ_x, σ_z by ones when needed
    σ_x = [σ <= 1e-10 ? 1.0 : σ for σ in σ_x]
    σ_z = [σ <= 1e-10 ? 1.0 : σ for σ in σ_z]

    # Build New Series Model
    new_model = SeriesModel([AffineMap(1./(nσ*σ_x) , -μ_x./(nσ*σ_x)), mdl,
                             AffineMap((nσ*σ_z), μ_z)])

    return new_model
end


using BioDataExpression
@static if VERSION < v"0.7.0-DEV.2005"
    using Base.Test
else
    using Test
end

# write your own tests here
@test 1 == 1

dbfile = joinpath(pwd(),"test", "data","database.sqlite");

@test isfile(dbfile)

data = BioData( dbfile );

@test length( nonzeros(data.data) ) == 97784

Normalize(data, 100)

@test length( nonzeros(data.raw) ) == 97784

@test length( nonzeros(data.data) ) == 77660

@test corspearman( log(toPlot( data.data[1:100,1])+1),log10(toPlot(data.raw[1:100,1])+1) ) > 0.9

res = zersos( size(data.data)[2] )
for i=1:size(data.data)[2]
  res[i] = corspearman( log(toPlot( data.data[1:100,i])+1),log10(toPlot(data.raw[1:100,i])+1) ) > 0.8
end

@test length(find(res .== 0)) == 0

Zscore(data)

@test length( nonzeros(data.zscored) ) == 77660

@test corspearman( log(toPlot( data.data[1:100,1])+1),log10(toPlot(data.zscored[1:100,1])+1) ) > 0.9

res = zersos( size(data.data)[2] )
for i=1:size(data.data)[2]
  res[i] = corspearman( log(toPlot( data.data[1:100,i])+1),log10(toPlot(data.zscored[1:100,i])+1) ) > 0.8
end

@test length(find(res .== 0)) == 0


@doc doc"""
    CircleGroup <: GroupManifold{Circle{ℂ},MultiplicationOperation}

The circle group is the complex circle ([`Circle(ℂ)`](@ref)) equipped with
the group operation of complex multiplication ([`MultiplicationOperation`](@ref)).
"""
const CircleGroup = GroupManifold{Circle{ℂ},MultiplicationOperation}

CircleGroup() = GroupManifold(Circle{ℂ}(), MultiplicationOperation())

show(io::IO, ::CircleGroup) = print(io, "CircleGroup()")

function compose(G::CircleGroup, x::AbstractVector, y::AbstractVector)
    return map(compose, repeated(G), x, y)
end

compose!(G::CircleGroup, z, x, y) = copyto!(z, compose(G, x, y))

identity(::CircleGroup, x::AbstractVector) = map(one, x)
identity(G::GT, e::Identity{GT}) where {GT<:CircleGroup} = e

identity!(::CircleGroup, y::AbstractVector, x) = copyto!(y, 1)
identity!(::GT, y::AbstractVector, ::Identity{GT}) where {GT<:CircleGroup} = copyto!(y, 1)

inv(G::CircleGroup, x::AbstractVector) = map(inv, repeated(G), x)
inv(G::GT, e::Identity{GT}) where {GT<:CircleGroup} = e

function inverse_translate(
    ::CircleGroup,
    x::AbstractVector,
    y::AbstractVector,
    ::LeftAction,
)
    return map(/, y, x)
end
function inverse_translate(
    ::CircleGroup,
    x::AbstractVector,
    y::AbstractVector,
    ::RightAction,
)
    return map(/, y, x)
end

translate_diff(::GT, x, y, v, ::ActionDirection) where {GT<:CircleGroup} = map(*, x, v)
function translate_diff(
    ::GT,
    ::Identity{GT},
    y,
    v,
    ::ActionDirection,
) where {GT<:CircleGroup}
    return v
end

function translate_diff!(G::CircleGroup, vout, x, y, v, conv::ActionDirection)
    return copyto!(vout, translate_diff(G, x, y, v, conv))
end


using Lasso
DATADIR = joinpath(dirname(@__FILE__), "data")

# Test against fused lasso results from glmgen
lakehuron = readcsvmat(joinpath(DATADIR, "LakeHuron.csv");header=true)[:, 3]
@testset "FusedLasso" begin
	@testset "λ = $(lambda)" for lambda in (10, 1, 0.1)
		@test coef(fit(FusedLasso, lakehuron, lambda)) ≈ vec(readcsvmat(joinpath(DATADIR, "LakeHuron_lambda_$lambda.csv")))
	end
end


module fresnelltools

using LinearAlgebra

using ..matrixcore

export FresnellBoundrary, FresnellSlab, Grating, ContinousBorder, ThreeLayerSystem, Interrogator, BorderInterrogator

function FresnellBoundrary(n1::Number, n2::Number)::Matrix{Number}
    # Equation from Saleh & Teich 3.ed. eq.6.2-8 & eq. 6.2-9
    # Returns the scattering matrix for the relevant border

    FresnellBoundrary(n1, n2, 0)[1]
end

function FresnellBoundrary(n1::Number, n2::Number, θ1::Number)::Tuple{Matrix{Number}, Matrix{Number}, Number}
    # Equation from Saleh & Teich 3.ed. eq.6.2-8 & eq. 6.2-9
    # It is similar to the other FresnellBoundrary function, but takes incident angle into account
    # Returns the TE scattering matrix, the TM scattering matrix and the outgoing transmitted angle

    # The imaginary coefficient 1e-14im has no physical reason to exist, and is only added for stability in the program.
    # It is assumed to be needed because of type instability in the asin function i Julia.
    # See https://github.com/JuliaLang/julia/issues/24296
    if ((typeof(n1) <: Complex) & (typeof(n2) <: Real))
        complex_modifier = 1e-14im
    else
        complex_modifier = 0im
    end
    θ2 = asin(complex_modifier + sin(θ1) * n1 / n2)

    r_te_12 = (n1 * cos(θ1) - n2 * cos(θ2)) / (n1 * cos(θ1) + n2 * cos(θ2))
    r_te_21 = (n2 * cos(θ2) - n1 * cos(θ1)) / (n2 * cos(θ2) + n1 * cos(θ1))

    r_tm_12 = (n1 * sec(θ1) - n2 * sec(θ2)) / (n1 * sec(θ1) + n2 * sec(θ2))
    r_tm_21 = (n2 * sec(θ2) - n1 * sec(θ1)) / (n2 * sec(θ2) + n1 * sec(θ1))

    t_te_12 = 1 + r_te_12
    t_te_21 = 1 + r_te_21

    t_tm_12 = (1 + r_tm_12) * cos(θ1) / cos(θ2)
    t_tm_21 = (1 + r_tm_21) * cos(θ2) / cos(θ1)

    Ste = [
        t_te_12 r_te_21;
        r_te_12 t_te_21
    ]
    Stm = [
        t_tm_12 r_tm_21;
        r_tm_12 t_tm_21
    ]

    return Ste, Stm, θ2
end

function FresnellSlab(n::Number, k_0::Number, d::Number, θ::Number)::Matrix{Number}
    # Equation from Saleh & Teich 3.ed. eq.7.1-4
    delay = ℯ^(-1im * n * k_0 * d * cos(θ))
    S = [
        delay 0;
        0     delay
    ]

    if abs(S[1, 1]) < 1e-5
        S /= abs(S[1, 1])
        S *= 1e-5
    end

    S
end

function Grating(n_spechial::Number, n_normal::Number, d_spechial::Real, d_normal::Real, layers::Int, k_0::Number)::Matrix{Number}
    # A utility for creating the scattering matrix of a periodic grating

    border_normal_spechial = FresnellBoundrary(n_normal, n_spechial)
    bulk_spechial = FresnellSlab(n_spechial, k_0, d_spechial, 0)
    border_spechial_normal = FresnellBoundrary(n_spechial, n_normal)
    bulk_normal = FresnellSlab(n_normal, k_0, d_normal, 0)

    cell = CascadeScattering([border_normal_spechial, bulk_spechial, border_spechial_normal, bulk_normal])
    Repeating(cell, layers)
end

function Grating(n_spechial::Number, n_normal::Number, d_spechial::Real, d_normal::Real, layers::Int, k_0::Number, θ1::Real)::Tuple{Matrix{Number}, Matrix{Number}}
    # A utility similar to Grating, but it takes an angle and assumes an infinitely wide grating
    # Returns the TE and TM scattering matrices for the grating
    # Unlike the angleded FresnellBoundrary tool, no angle is returned, this is because the grating assumes normal material on both sides

    n_s_te, n_s_tm, θ2 = FresnellBoundrary(n_normal, n_spechial, θ1)
    s_n_te, s_n_tm, _  = FresnellBoundrary(n_spechial, n_normal, θ2)

    n_b = FresnellSlab(n_normal, k_0, d_normal, θ1)
    s_b = FresnellSlab(n_spechial, k_0, d_spechial, θ2)

    cell_te = CascadeScattering([n_s_te, s_b, s_n_te, n_b])
    cell_tm = CascadeScattering([n_s_tm, s_b, s_n_tm, n_b])

    Repeating(cell_te, layers), Repeating(cell_tm, layers)
end

function ContinousBorder(n_from::Number, n_to::Number, d::Real, k_0::Number, stepps::Int)::Matrix{Number}
    # A utility for creating a semi-continous change in refractive index
    # The gradient is approximated as linear

    components = Array{Matrix{Number}}(undef, 2 * stepps)

    n_stepp = (n_to - n_from) / stepps
    d_stepp = d / stepps
    for i in 1:stepps
        n1 = n_from + (i - 1) * n_stepp
        n2 = n1 + n_stepp

        components[2 * i - 1] = FresnellBoundrary(n1, n2)
        components[2 * i] = FresnellSlab(n2, k_0, d_stepp, 0)
    end

    CascadeScattering(components)
end

function ThreeLayerSystem(n_1::Number, n_bulk, n_3::Number, λ, θ, d)
    interface1_te, interface1_tm, θ2 = FresnellBoundrary(n_1, n_bulk(λ), θ)
    bulk = FresnellSlab(n_bulk(λ), 2*π / λ, d, θ2)
    interface2_te, interface2_tm, _ = FresnellBoundrary(n_bulk(λ), n_3, θ2)
    te_system = CascadeScattering([interface1_te, bulk, interface2_te])
    tm_system = CascadeScattering([interface1_tm, bulk, interface2_tm])

    te_system, tm_system
end

function Interrogator(layers::Vector{Function}, distances::Vector{Float64}, step::Float64)::Tuple{Vector{Float64}, Vector{Number}}
    # A small utility providing a default lambda for the Interrogator tool
    Power(Up, Um) = abs(Up + Um)^2
    Interrogator(layers, distances, step, Power)
end

function Interrogator(layers::Vector{Function}, distances::Vector{Float64}, step::Float64, expression::Function)::Tuple{Vector{Float64}, Vector{Number}}
    # A tool for inspecting the field inside a multilayer system

    total_system = Array{Matrix{Number}}(undef, length(layers))
    for i in 1:1:length(layers)
        total_system[i] = layers[i](distances[i])
    end
    U0_plus = 1
    U0_minus = (CascadeScattering(total_system) * [U0_plus, 0])[2]

    maxlen = ceil(Int, length(distances) * 2 + sum(distances) / step)
    d_values = Vector{Number}(undef, maxlen)
    u_values = Vector{Number}(undef, maxlen)

    d_values[1] = 0.0
    u_values[1] = expression(U0_plus, U0_minus)

    m_accumulated = I
    m_local = I
    m_partial = I
    d_accumulated = 0.0
    i = 1
    for (layer, distance) in zip(layers, distances)
        m_local = m_accumulated

        m_partial = StoM(layer(step))
        j = 0
        for _ in 0:step:distance
            i += 1
            j += 1
            d_accumulated += distance > step ? step : distance
            if (j % 1000) == 0
                m_local = StoM(layer(step * j)) * m_accumulated
            else
                m_local = m_partial * m_local
            end

            U_plus, U_minus = m_local * [U0_plus, U0_minus]
            measure = expression(U_plus, U_minus)
            d_values[i] = d_accumulated
            u_values[i] = measure
        end

        d_rest = distance % step
        if (d_rest > (step / 1000)) & (step > distance)
            m_partial = StoM(layer(d_rest))
            i += 1
            d_accumulated += d_rest
            m_local = m_partial * m_local

            (U_plus, U_minus) = m_local * [U0_plus, U0_minus]

            measure = expression(U_plus, U_minus)
            d_values[i] = d_accumulated
            u_values[i] = measure
        end

        m_accumulated = StoM(layer(distance)) * m_accumulated
    end

    u_values[1:i], d_values[1:i]
end

function BorderInterrogator(layers::Vector{Matrix{Number}}, distances::Vector{Float64}, expression::Function)::Tuple{Vector{Number}, Vector{Number}, Vector{Real}}
    # A tool for calculating the field at the borders in a structure independent of the Interrogator function
    # Returns the forward propagating fields, the backward propagating fields and the distances

    U0p = 1
    U0n = (CascadeScattering(layers) * [U0p, 0])[2]

    Upv = Vector{Number}()
    Unv = Vector{Number}()
    Dv = Vector{Float64}()

    push!(Upv, U0p)
    push!(Unv, U0n)
    push!(Dv, 0.0)

    accumulated_m = I
    accumulated_d = 0.0
    for (layer, distance) in zip(layers, distances)
        accumulated_m = StoM(layer) * accumulated_m
        accumulated_d += distance

        Up, Un = accumulated_m * [U0p, U0n]

        push!(Upv, Up)
        push!(Unv, Un)
        push!(Dv, accumulated_d)
    end

    return expression.(Upv), expression.(Unv), Dv
end

end # fresnelltools


include("nanop.jl")

using NCDatasets
using Formatting
using Statistics

Re  = 6371000.0 #  m
c_p  = 1004.0   #  J / kg K
g   = 9.81      #  m / s^2
Lw  = 2.26e6    #  J / kg

linestyle = ["r-", "g-", "b-"]

Dataset("domain_nc_files/domain.lnd.fv4x5_gx3v7.091218.nc", "r") do ds
    global lat, lon, lev, ilev, rec

    lat = replace(ds["yc"][1,:], missing=>NaN)
    lon = replace(ds["xc"][:,1], missing=>NaN)

end

using PyPlot
plt[:rcParams]["font.size"] = 40
plt[:rcParams]["axes.label.size"] = 40

fig, ax = plt[:subplots](1, 1, figsize=(12,8))
#ax[:plot]([-90, 90], [0, 0], linewidth=2, color="#cccccc")

casename = "lowres_SSM_NK"

Dataset("extract_nc_files/$casename.ocn.h.ma.fv4x5.0001-0020.nc", "r") do ds

    t_beg = 240 - 10*12 + 1
    t_end = 240

    cnt_JJA = 0
    cnt_DJF = 0

    mld_JJA = zeros(length(lon), length(lat))
    mld_DJF = zeros(length(lon), length(lat))

    for t = t_beg:t_end
        m = mod(t - 1, 12) + 1

        if m in [11, 12, 1]
            mld_DJF += replace(ds["mld"][:, :, t], missing=>NaN)
            cnt_DJF += 1
        elseif m in [6, 7, 8]
            mld_JJA += replace(ds["mld"][:, :, t], missing=>NaN)
            cnt_JJA += 1
        end

    end

    mld_JJA ./= cnt_JJA
    mld_DJF ./= cnt_DJF

    mld_JJA = nanmean(mld_JJA, dims=(1,))[1, :]
    mld_DJF = nanmean(mld_DJF, dims=(1,))[1, :]

    ax[:plot](lat, mld_DJF, "k-", linewidth=2,  label="DJF")
    ax[:plot](lat, mld_JJA, "k--", linewidth=2,  label="JJA")

end

ax[:legend](loc="lower center", fontsize=15)
ax[:grid](linestyle="--")
ax[:set_xlabel]("Latitude [deg]", size=15)
ax[:set_ylabel]("Mixed-layer Depth [m]", size=15)
ax[:set_xlim](-90, 90)
ax[:set_ylim](-5, 800)

ax[:invert_yaxis]()

ax[:set_title]("lowres_SSM_NK 11-20 year MLD average in DJF(solid) and JJA(dashed).")

fig[:savefig]("$(casename)_MLD.png", dpi=200)


using LinearAlgebra, ForwardDiff, Optim
using DataFrames, Plots, Distributions, Statistics

ProjDir = @__DIR__
cd(ProjDir)

include("./dynamic_pricing/price_demand_models.jl")
include("./demand_curve/select_optimal_demand_curve.jl")
include("./intervals/intervals.jl")
include("./utils/logx.jl")
include("./actual_demand/actual_demand_curve.jl")

# Generate demand hypothesis

a_range = range(-1.2, -0.6, length=4)
b_range = range(60, 70, length=4)
h_vec = generate_demand_hypothesis(a_range, b_range)
println("Price demand hypotheses (see also fig_01.png):")
display(h_vec)
println()

plot(xlim=(0, 60), ylim=(0, 80),
    xlabel="Price", ylabel="Demand",
    title="Price Demand hypotheses with optimal price points."
)
x = range(0, 60, length=101)
for i in 1:size(h_vec, 1)
	plot!(x, h_vec[i, :d].(x), color=:red, leg=false)
end
scatter!(h_vec[:, :p_opt], h_vec[:, :d_opt], m=(:cross))
savefig(joinpath(ProjDir, "fig_01.png"))

# Emperical demand curve

println()
curve = emperical_demand_curve(30, 60, 20)
println("Emperical price demand curve:")
display(curve)

# Intervals for price updates

T = 24 * 1                       # time step is one hour, flash offering for 1 day
m = 4                            # not more than 4 price updates

t_mask = intervals(m, T, 2)      # intervals for price lvels
p = h_vec[1, :p_opt]             # initial price

println("\nSelected intervals:")
display(t_mask')
println("\nInitial price = $p\n")

# Do a first iteration in selecting demand hypothesis

(price, row, hist_d) = select_optimal_demand_curve(t_mask, p)
println()
display((price, row))
println()
display((hist_d))

plot(xlim=(0, 60), ylim=(0, 80), xlabel="Price", ylabel="Demand")
plot!(curve[:, :prices], curve[:, :demands],
    xlabel="price", ylabel="demand", line=(:dot, :gray))
x = range(0, 60, length=101)
plot!(x, h_vec[row, :d].(x), color=:red, leg=false)
savefig(joinpath(ProjDir, "fig_02.png"))

#=
colors = ['#8EA604', '#F5BB00', '#EC9F05', '#D76A03']

def visualize_snapshot(t):
    fig.clear()

    plt.subplot(2, 1, 1)
    plt.xlabel('Price')
    plt.ylabel('Demand')
    plt.title('Price-demand curve fitting')
    plt.plot(curve[0], curve[1], 'k:')
    for i in range(0, t):
        plt.plot(prices, list(map(np.array(history)[i, 2], prices)), 'k-')
    plt.plot(np.array(history)[0:t,0], np.array(history)[0:t, 1], 'rx')

    plt.subplot(2, 1, 2)
    plt.xlabel('Time')
    plt.ylabel('Demand/price')
    plt.title('Realized demand and price')
    plt.plot(range(0, t+1), np.array(history)[0:t+1,0], 'r-')
    bars = plt.bar(range(0, T-1), np.pad(np.array(history)[0:t+1, 1], (0, T-2-t), 'constant'), 0.35)
    for i in range(0, t+1):
        bars[i].set_color(colors[t_mask[i]])

fig = plt.figure(figsize=(10, 8))
plt.subplots_adjust(hspace=0.5)
visualize_snapshot(T-2)               # visualize the state in the end of the simulation
plt.show()
=#


# ---
# title: 765. Couples Holding Hands
# id: problem765
# author: Indigo
# date: 2021-02-14
# difficulty: Hard
# categories: Greedy, Union Find, Graph
# link: <https://leetcode.com/problems/couples-holding-hands/description/>
# hidden: true
# ---
#
# N couples sit in 2N seats arranged in a row and want to hold hands. We want to
# know the minimum number of swaps so that every couple is sitting side by side.
# A _swap_ consists of choosing **any** two people, then they stand up and
# switch seats.
#
# The people and seats are represented by an integer from `0` to `2N-1`, the
# couples are numbered in order, the first couple being `(0, 1)`, the second
# couple being `(2, 3)`, and so on with the last couple being `(2N-2, 2N-1)`.
#
# The couples' initial seating is given by `row[i]` being the value of the
# person who is initially sitting in the i-th seat.
#
# **Example 1:**
#
#
#
#     Input: row = [0, 2, 1, 3]
#     Output: 1
#     Explanation: We only need to swap the second (row[1]) and third (row[2]) person.
#
#
# **Example 2:**
#
#
#
#     Input: row = [3, 2, 0, 1]
#     Output: 0
#     Explanation: All couples are already seated side by side.
#
#
# **Note:**
#
#   1. `len(row)` is even and in the range of `[4, 60]`.
#   2. `row` is guaranteed to be a permutation of `0...len(row)-1`.
#
#
## @lc code=start
using LeetCode

function min_swaps_couples(row::Vector{Int})
    len, ps, res = length(row), length(row) ÷ 2, 0
    grp = [Set{Int}() for i in 1:ps]
    for i in 1:2:len
        l, r = row[i] ÷ 2 + 1, row[i + 1] ÷ 2 + 1
        (l != r) && (push!(grp[l], r); push!(grp[r], l))
    end
    visited = fill(false, ps)
    q = Queue{Int}()
    for i in 1:ps
        visited[i] && continue
        enqueue!(q, i)
        visited[i] = true
        cnt = 0
        while !isempty(q)
            frt = dequeue!(q)
            cnt += 1
            for neighbor in grp[frt]
                visited[neighbor] && continue
                enqueue!(q, neighbor)
                visited[neighbor] = true
            end
        end
        res += cnt - 1
    end
    res
end
## @lc code=end


# This is adapted from Base.isapprox
# https://github.com/JuliaLang/julia/blob/381693d3dfc9b7072707f6d544f82f6637fc5e7c/base/floatfuncs.jl#L222-L291
function isapprox_le(x::Number, y::Number; atol::Real=0, rtol::Real=Base.rtoldefault(x,y,atol))
    x <= y || (isfinite(x) && isfinite(y) && abs(x-y) <= max(atol, rtol*max(abs(x), abs(y))))
end


export rlocus

"""
### function rlocus(G::TransferFunction, k = []; plot = true)

Compute and, optionally, plot the root locus of the open-loop system represented
by the transfer function `G`.

##### Args

* G: Transfer function.
* k: Vector with the open-loop gains to compute the root locus.

##### Parameters

* plot: If true, then the root locus will be plotted.

##### Returns

* root_locus: An array with the closed-loop poles of the system.
* root_locus_k: An array with the gains used to compute the closed-loop gains.

"""

function rlocus(G::TransferFunction, k = []; plot = true)
    # Check the transfer function.
    if length(G.num.a) > length(G.den.a)
        error("rlocus() does not support improper systems.")
    end

    # Compute the poles and zeros of the transfer function.
    p = convert(Array{Complex128,1}, pole(G))
    sort!(p, by=real)

    z = convert(Array{Complex128,1}, izero(G))

    # Minimal realization of the transfer function.
    Gm = minreal(G)
    mp = pole(Gm)
    mz = izero(Gm)
    sort!(mp, by=real)

    # Array to store the root locus.
    root_locus   = Array{Array{Complex128,1}}(0)
    root_locus_k = Array{Float64,1}(0)

    # Select gains to compute the root locus.
    # =======================================
    #
    # The idea of some parts of this algorithm was obtained from the file
    # `rlocus.m` from Octave-Forge.

    k_user = true

    if isempty(k)
        # Indicate that the use did not provided a gain vector.
        k_user = false

        # In the following, there is an ad-hoc algorithm to estimate the highest
        # gain necessary to provide the entire information to the user.

        # Real-axis break-in and breakaway points.
        # ========================================

        # Compute the real-axis break-in and breakaway points.
        num_dKds = polyder(Gm.den)*Gm.num - Gm.den*polyder(Gm.num)

        # Compute the points.
        break_points = roots(num_dKds)

        # Compute the gains related to those points.
        k_crossing =
            real(map(s->polyval(-G.den,s)/polyval(G.num, s), break_points))

        # Get the maximum gain. If it is empty, then use 100.
        k_max = maximum( [k_crossing; 100.0] )

        # Asymptotes
        # ==========

        k_asym = 0.0

        # Number of asymptotes.
        num_asymptotes = length(mp) - length(mz)

        if num_asymptotes > 0
            # Angle and start point for the asymptotes.
            σ_a = (sum(mp) - sum(mz))/num_asymptotes
            θ_a = [ (2*m+1)*pi/num_asymptotes for m in 0:(num_asymptotes-1) ]

            # Get 2x the maximum distance from the center of the asymptotes to
            # the open-loop poles and zeros and the break points.
            aux = [mp; mz; break_points]

            Δ_max  = 0.0

            if !isempty(aux)
                Δ_max = 2*maximum(abs.( [mp; mz; break_points] - σ_a ) )
            end

            # Get the gain related to Δ_max for each asymptotes.
            s = σ_a + Δ_max*exp.(complex.(0.0, θ_a))
            k_asym = real(map(s->polyval(-G.den,s)/polyval(G.num, s), s))
        end

        # Maximum gain
        # ============

        # Update the maximum gain given the previous information.
        k_max = maximum( [k_max; k_asym] )

        # Compute 50 points between 0 and k_max.
        k = collect(linspace(0, k_max, 50))
    end

    # Compute the tolerance for smoothing.
    # ====================================

    # Compute the tolerance for smoothing the root locus.
    ϵ = 0.0

    if !isempty(mp)
        η = [abs.( maximum(real(mp)) - minimum(real(mp)) );
             abs.( maximum(imag(mp)) - minimum(imag(mp)) ); ]

        # Julia function `minimum` fails if the argument is empty.
        (!isempty(η)) && (ϵ = 0.01*minimum(η))
    end

    # The tolerance must not be 0.
    (ϵ == 0.0) && (ϵ = 1e-2)

    # Compute the root locus.
    # =======================

    # Get the poles of the minimal realization of the system.
    push!(root_locus,   mp)
    push!(root_locus_k, 0.0)

    # Compute closed-loop poles.
    pole_k_1 = mp

    i = 1
    ki = k[i]

    # Notice that, if the minimum realization of the system has no poles, then
    # nothing must be computed.
    while !isempty(mp) && (i <= length(k))
        ki = k[i]

        (ki == 0) && (i += 1; continue)

        # Closed-loop dynamic matrix.
        den_cl = Gm.den + ki*Gm.num

        pole_k = convert(Array{Complex128,1},roots(den_cl))

        # Sort the poles for ki to be as closest as possible to the poles
        # computed in the last iteration.
        for j = 1:length(pole_k_1)
            # Do not consider the past values, since they are already sorted.
            ind = indmin(abs.(pole_k[j:end]-pole_k_1[j]))

            aux             = pole_k[j]
            pole_k[j]       = pole_k[j+ind-1]
            pole_k[j+ind-1] = aux
        end

        # If the user did not provide a gain vector, then add more points to
        # smooth the root locus.
        if !k_user
            # Compute the maximum distance between to current poles and the
            # poles of the last iteration.
            Δ = maximum(norm.(pole_k[1:length(pole_k_1)] - pole_k_1))

            # If the difference if too high, add more gains.
            if (Δ > ϵ) && (k[i-1] < k[i])
                k_new = collect(linspace(k[i-1], k[i], 10))

                if i < length(k)
                    k = [k[1:i-2]; k_new; k[i+1:end]]
                else
                    k = [k[1:i-2]; k_new;]
                end

                continue
            end
        end

        # Poles of the closed-loop system. Notice that we should only select the
        # number of poles in `pole_k_1`. The other should be those that will be
        # cancelled with the zeros.
        pole_k_1 = copy(pole_k[1:length(pole_k_1)])
        push!(root_locus, pole_k_1)
        push!(root_locus_k, ki)
        i += 1
    end

    # Now, put the zeros into the root locus by comparing the distance between
    # them and the latest computed closed-loop poles.
    aux      = copy(pole_k_1)
    zeros_k  = NaN*ones(pole_k_1)

    for i = 1:length(mz)
        ind = indmin(abs.(mz[i]-aux))
        zeros_k[ind] = mz[i]

        # Remove the matched pole to avoid problems with zeros that have
        # multiplicity higher than 1.
        aux[ind] = Inf
    end

    push!(root_locus, zeros_k)
    push!(root_locus_k, Inf)

    # Plot the root locus.
    # ====================

    # Check if the root locus must be plotted.
    if plot
        mz = izero(Gm)
        plot_rlocus(root_locus, root_locus_k, p, z, mp, mz)
    end

    root_locus, root_locus_k
end


### A Pluto.jl notebook ###
# v0.17.5

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el)
        el
    end
end

# ╔═╡ 08980845-f030-4a64-a4b2-ab027b3a2721
   begin  using Pkg
       Pkg.activate(".testenv")
	   Pkg.add("Revise")
       using Revise
       Pkg.add(["Triangulate","ExtendableGrids"])
	   Pkg.add("GridVisualize")
	   Pkg.add("PlutoUI")
	   Pkg.develop("PlutoVista")
   end

# ╔═╡ d6c0fb79-4129-444a-978a-bd2222b53df6
begin
    using PlutoUI
    using PlutoVista
    using GridVisualize
    using Printf
    using Triangulate
	using ExtendableGrids
end

# ╔═╡ 93ca4fd0-8f61-4174-b459-55f5395c0f56
md"""
# Test Notebook for [PlutoVista](https://github.com/j-fu/PlutoVista.jl) and [vtk.js](https://kitware.github.io/vtk-js/index.html)
"""

# ╔═╡ 7c06fcf0-8c98-49f7-add8-435f57a9c9da
function maketriangulation(maxarea)
    triin=Triangulate.TriangulateIO()
    triin.pointlist=Matrix{Cdouble}([-1.0 -1.0 ; 1.0 -1.0 ; 1.0 2 ; -1.0 1.0]')
    triin.segmentlist=Matrix{Cint}([1 2 ; 2 3 ; 3 4 ; 4 1; 4 2 ]')
    triin.segmentmarkerlist=Vector{Int32}([1, 2, 3, 4, 5])
    triin.regionlist=Matrix{Cdouble}([-0.9 0.9; -0.9 0.9; 1 2 ; 2.0*maxarea maxarea])
    (triout, vorout)=triangulate("paADQ", triin)
    triout.pointlist, triout.trianglelist,Int.(vec(triout.triangleattributelist)),triout.segmentlist,triout.segmentmarkerlist
end

# ╔═╡ 60dcfcf5-391e-418f-8e7c-3a0fe94f1e0d
p=PlutoVTKPlot(resolution=(300,300),tickfontsize=10)

# ╔═╡ db2823d9-aa6d-4be3-af5c-873c072cfd2b
md"""
Change grid resolution: $(@bind resolution Slider(5:200))
"""

# ╔═╡ 890710fe-dac0-4256-b1ba-79776f1ea7e5
(pts,tris,markers,edges,edgemarkers)=maketriangulation(1/resolution^2)

# ╔═╡ b8a976e3-7fef-4527-ae6a-4da31c93a04f
func=0.5.+0.5*[sin(10*pts[1,i])*cos(10*pts[2,i]) for i=1:size(pts,2)]

# ╔═╡ 401b36bd-fa8f-4a9c-9556-bbc82c3ddbca
 md"""
Change time: $(@bind time Slider(0:0.1:10,show_value=true))
"""

# ╔═╡ 6fd4a1ee-7a4a-405b-8e1f-5819eababe10
ft=0.5*[sin(10*pts[1,i]-time)*cos(10*pts[2,i]-time) for i=1:size(pts,2)];

# ╔═╡ e76f8a6a-ab91-454a-b200-cfc8b57eb331
triplot!(p,pts,tris,ft)

# ╔═╡ bce0cfe7-4112-4bb8-aac6-43885f3746a9
md"""Number of gridpoints: $(size(pts,2)) """

# ╔═╡ 81046dcd-3cfb-4133-943f-61b9b3cdb183
let
	p=PlutoVTKPlot(resolution=(300,300),axisfontsize=20)
	tricontour!(p,pts,tris,func;cmap=:spring,levels=(0.1:0.2:1),limits=(0,1),xlabel="a",aspect=1.5)
end

# ╔═╡ 7019ce3f-f2db-4581-8bd9-64f76231a62a
let
	p=PlutoVTKPlot(resolution=(300,300),legendfontsize=15)
	trimesh!(p,pts,tris;
		markers=markers,
		edges=edges,edgemarkers=edgemarkers
	)
end

# ╔═╡ 2e3546f6-eb47-4693-aa00-902570fab7b5
function grid3d(;n=15)
    X=collect(0:1/n:1)
    g=simplexgrid(X,X,X)
end

# ╔═╡ bd0a59a2-564d-42bd-ab6f-a50b26f9241f
function func3d(;n=15)
    g=grid3d(n=n)
    g, map((x,y,z)->sinpi(2*x)*sinpi(3.5*y)*sinpi(1.5*z),g)
end

# ╔═╡ 368b8cf5-fabd-4b84-b33c-b15c4452393b
	g,f=func3d(;n=20)

# ╔═╡ f64729e4-d2b4-40d3-acbb-1395dbe0337d
p3d=PlutoVTKPlot(resolution=(300,300))

# ╔═╡ 7be35f33-5f7a-4765-8bb6-1487e209efc8
md"""
f: $(@bind flevel Slider(0:0.01:1,show_value=true,default=0.45))

x: $(@bind xplane Slider(0:0.01:1,show_value=true,default=0.45))
y: $(@bind yplane Slider(0:0.01:1,show_value=true,default=0.45))
z: $(@bind zplane Slider(0:0.01:1,show_value=true,default=0.45))
"""

# ╔═╡ 3681ef5b-c794-44da-9fe7-cedcd68b426c
tetcontour!(p3d,g[Coordinates],g[CellNodes],f;levels=[flevel],
	faces=g[BFaceNodes],
	facemarkers=g[BFaceRegions],
	xplanes=[xplane],yplanes=[yplane],zplanes=[zplane],outlinealpha=0.1, levelalpha=0.25)

# ╔═╡ 0f440c27-7ff1-4db5-b4eb-8ce1e9018ef1
g[BFaceNodes]

# ╔═╡ 606f6837-f3b7-4a52-b9c3-034799c7bf93
g3x=let
	n=10
X=collect(0:1/n:1)
simplexgrid(X,X,X)
end

# ╔═╡ ecb3bb5e-6ae5-4d6e-9834-d52ce977b3fc
p3dx=PlutoVTKPlot(resolution=(300,300))

# ╔═╡ 90ff6ffc-84dc-45fd-8d09-9eb916397630
 md"""
x: $(@bind gxplane Slider(0:0.01:1,show_value=true,default=0.45))
y: $(@bind gyplane Slider(0:0.01:1,show_value=true,default=0.45))
z: $(@bind gzplane Slider(0:0.01:1,show_value=true,default=0.45))
"""

# ╔═╡ ae5707d9-41b7-4937-924a-fb54b83c31db
alpha=0.15

# ╔═╡ 519d106f-3f6f-4db1-b4e4-c5e7ef176857
tetmesh!(p3dx,g3x[Coordinates],g3x[CellNodes];
		markers=g3x[CellRegions],
	faces=g3x[BFaceNodes],
	facemarkers=g3x[BFaceRegions],
	xplanes=gxplane,yplanes=gyplane,zplanes=gzplane,outlinealpha=0)

# ╔═╡ 4263e897-d878-4fab-acae-a6c4dae37c5e
qp=PlutoVTKPlot(resolution=(500,500))

# ╔═╡ 5dfb1dde-06ed-483a-bac3-6a21a7f98856
@bind x0 Slider(1:0.01:9)

# ╔═╡ 6d7547c6-5667-463f-99dc-2f3fbaeb4c4d
let
	n=50
	X=0:10/n:10
	g=simplexgrid(X,X)
	pts=g[Coordinates]
    f(x,y)=sin(x-x0)*cos(y)+0.05*(x-x0)*y
    fx(x,y)=-cos(x-x0)*cos(y)-0.05*y
    fy(x,y)=sin(x-x0)*sin(y)-0.05*(x-x0)
    v=map(f,g)
    ∇v=0.5*vcat(map(fx,g)',map(fy,g)')

	tricontour!(qp,pts,g[CellNodes],v,colormap=:summer,levels=5)
	quiver2d!(qp,pts,∇v)
end

# ╔═╡ 23c75823-e69f-4ca0-a3ca-783612ba9c3c
html"""<hr>"""

# ╔═╡ Cell order:
# ╟─93ca4fd0-8f61-4174-b459-55f5395c0f56
# ╠═d6c0fb79-4129-444a-978a-bd2222b53df6
# ╠═7c06fcf0-8c98-49f7-add8-435f57a9c9da
# ╠═890710fe-dac0-4256-b1ba-79776f1ea7e5
# ╠═b8a976e3-7fef-4527-ae6a-4da31c93a04f
# ╠═60dcfcf5-391e-418f-8e7c-3a0fe94f1e0d
# ╟─db2823d9-aa6d-4be3-af5c-873c072cfd2b
# ╟─401b36bd-fa8f-4a9c-9556-bbc82c3ddbca
# ╠═6fd4a1ee-7a4a-405b-8e1f-5819eababe10
# ╠═e76f8a6a-ab91-454a-b200-cfc8b57eb331
# ╟─bce0cfe7-4112-4bb8-aac6-43885f3746a9
# ╠═81046dcd-3cfb-4133-943f-61b9b3cdb183
# ╠═7019ce3f-f2db-4581-8bd9-64f76231a62a
# ╠═2e3546f6-eb47-4693-aa00-902570fab7b5
# ╠═bd0a59a2-564d-42bd-ab6f-a50b26f9241f
# ╠═368b8cf5-fabd-4b84-b33c-b15c4452393b
# ╠═f64729e4-d2b4-40d3-acbb-1395dbe0337d
# ╟─7be35f33-5f7a-4765-8bb6-1487e209efc8
# ╠═3681ef5b-c794-44da-9fe7-cedcd68b426c
# ╠═0f440c27-7ff1-4db5-b4eb-8ce1e9018ef1
# ╠═606f6837-f3b7-4a52-b9c3-034799c7bf93
# ╠═ecb3bb5e-6ae5-4d6e-9834-d52ce977b3fc
# ╠═90ff6ffc-84dc-45fd-8d09-9eb916397630
# ╠═ae5707d9-41b7-4937-924a-fb54b83c31db
# ╠═519d106f-3f6f-4db1-b4e4-c5e7ef176857
# ╠═4263e897-d878-4fab-acae-a6c4dae37c5e
# ╠═5dfb1dde-06ed-483a-bac3-6a21a7f98856
# ╠═6d7547c6-5667-463f-99dc-2f3fbaeb4c4d
# ╟─23c75823-e69f-4ca0-a3ca-783612ba9c3c
# ╠═08980845-f030-4a64-a4b2-ab027b3a2721


using WeavePynb
using Mustache

mmd(fname) = mmd_to_html(fname, BRAND_HREF="../toc.html", BRAND_NAME="Calculus with Julia")
## uncomment to generate just .md files
#mmd(fname) = mmd_to_md(fname, BRAND_HREF="../toc.html", BRAND_NAME="Calculus with Julia")

fnames = ["polar_coordinates",
          "vectors",
          "vector_valued_functions",
          "scalar_functions",
          "scalar_functions_applications",
          "vector_fields"
]

function process_file(nm, twice=false)
    include("$nm.jl")
    mmd_to_md("$nm.mmd")
    markdownToHTML("$nm.md")
    twice && markdownToHTML("$nm.md")
end

process_files(twice=false) = [process_file(nm, twice) for nm in fnames]

"""
## TODO differential_vector_calcululs

### Add questions for scalar_function_applications
* Newton's method??
* optimization. Find least squares for perpendicular distance using the same 3 points...??

"""


using Pkg
cd(@__DIR__)
Pkg.activate("../"); Pkg.instantiate()

using Plots; gr()
using LaTeXStrings
using DataFrames
using CSV
using OMJulia
using OrdinaryDiffEq
using LinearAlgebra
using ControlSystems
using DataInterpolations
using BlackBoxOptim
using BenchmarkTools

# Plot path
figpath = "../figs/";
# Linewidths and styles
LW1 = 2.5
LW2 = 1.5
LS1 = :solid
LS2 = :dot
LS3 = :dash
LC1 = :blue
LC2 = :red
LC3 = :green
LC4 = :black
LA1 = 1.0
LA2 = 0.7
LA3 = 0.3
;

# Function for helicopter model
function helicopter_ode(dx,x,Ps,t,P0,u)
# P0 = (rc=rc,h=h,m=m,Jbc=Jbc,Js=Js,g=g,K_th_th=K_th_th,K_th_psi=K_th_psi,K_psi_th=K_psi_th,K_psi_psi=K_psi_psi,D_th=D_th,D_psi=D_psi);
    th = x[1]
    psi = x[2]
    w_th = x[3]
    w_psi = x[4]
    #
    M11 = Ps[4]*P0.Jbc + Ps[3]*P0.m*((Ps[1]*P0.rc)^2+(Ps[2]*P0.h)^2)
    M22 = Ps[4]*P0.Jbc*cos(th)^2 + Ps[3]*P0.m*(Ps[1]*P0.rc*cos(th)-Ps[2]*P0.h*sin(th))^2 + Ps[5]*P0.Js
    #
    dLdth = -w_psi^2*(((Ps[4]*P0.Jbc)+(Ps[3]*P0.m)*((Ps[1]*P0.rc)^2-(Ps[2]*P0.h)^2))*cos(th)*sin(th)-(Ps[1]*P0.rc)*(Ps[2]*P0.h)*(cos(th)^2-sin(th)^2)) - (Ps[3]*P0.m)*(Ps[6]*P0.g)*((Ps[1]*P0.rc)*cos(th) - (Ps[2]*P0.h)*sin(th))
    #
    dp_psi_dth = -2*(((Ps[4]*P0.Jbc)+(Ps[3]*P0.m)*((Ps[1]*P0.rc)^2-(Ps[2]*P0.h)^2))*cos(th)*sin(th) + (Ps[3]*P0.m)*(Ps[1]*P0.rc)*(Ps[2]*P0.h)*(cos(th)^2-sin(th)^2))*w_psi
    #
    F_th = Ps[7]*P0.K_th_th*u.th(t) - Ps[8]*P0.K_th_psi*u.psi(t) - P0.D_th*w_th
    F_psi = Ps[9]*P0.K_psi_th*u.th(t) - Ps[10]*P0.K_psi_psi*u.psi(t) - P0.D_psi*w_psi
    #
    dx[1] = w_th
    dx[2] = w_psi
    dx[3] = (dLdth + F_th)/M11
    dx[4] = (-dp_psi_dth*w_th + F_psi)/M22
end
#
# Converting from radians to degrees in plots
#
rad2deg(t,u) = (t,u*180/pi)
;

# Model parameters
rc = 1.5e-2 # m
h = 0.5e-2 # m
m = 0.5 # kg
Jbc = 15e-3 # kg.m2
Js = 5e-3 # kg.m2
g = 9.81 # m/s2
K_th_th = 55e-3 # Nm/V
K_th_psi = 5e-3 # Nm/V
K_psi_th = 0.15 # Nm/V
K_psi_psi = 0.2 # Nm/V
D_th = 1e-2 # Nm/(rad/s)
D_psi = 8e-2 # Nm/(rad/s)
#
P0 = (rc=rc,h=h,m=m,Jbc=Jbc,Js=Js,g=g,K_th_th=K_th_th,K_th_psi=K_th_psi,K_psi_th=K_psi_th,K_psi_psi=K_psi_psi,D_th=D_th,D_psi=D_psi)
Ps = fill(1.0,length(P0))
;

# Set-up for simulating helicopter
#
th0 = pi/2
psi0 = 0.0
w_th0 = -0.1
w_psi0 = 0.1
#
x0 = [th0,psi0,w_th0,w_psi0]
u_c = (th=t->0.0, psi=t->0.0)
#
helicopter! = (dx,x,Ps,t) -> helicopter_ode(dx,x,Ps,t,P0,u_c)
#
tspan = (0.,20)
prob = ODEProblem(helicopter!,x0,tspan,Ps)
#
sol = solve(prob,AutoTsit5(Rosenbrock23()),reltol=1e-6)
#
plot(sol,vars=(rad2deg,0,1),lw=LW1,lc=LC1,label=L"$\theta$ [deg]")
plot!(sol,vars=(rad2deg,0,2),lw=LW1,lc=LC2,label=L"$\psi$ [deg]")
plot!(xlim = tspan,xlabel=L"time $t$ [s]")
figname="Helicopter_theoretical_simulation.svg"
savefig(figpath*figname)

df=CSV.read("../data/Lab-Helicopter_Experimental-data.csv")
ENV["COLUMNS"]=100 # change the number of IJulia columns from default 80 to 100
df[1:5,:]

u_th = df[:,1]
u_psi = df[:,2]
theta = df[:,3]
psi = df[:,4]
dt = df[1,5]
tm = 0:dt:(length(psi)-1)*dt
;

data = [theta';psi']
Nd = length(tm);

plot(tm,theta*180/pi,lw=LW2,lc=LC1,label=L"\theta",legend=:topleft)
plot!(twinx(),tm,psi*180/pi,lw=LW2,lc=LC2,label=L"\psi")
plot!(title="Experimental helicopter outputs: measured angles", ylabel=L"angles [$^\circ$]", xlabel=L"time $t$ [s]", box=true)
figname="Helicopter_experimental_angles.svg"
savefig(figpath*figname)

plot(tm,u_th,lw=LW1,lc=LC1,label=L"u_\theta",legend=:topleft)
plot!(twinx(),tm,u_psi,lw=LW1,lc=LC2,label=L"u_\psi",legend=:bottomright)
plot!(title="Experimental helicopter inputs: motor voltages", ylabel=L"$u$ [V]", xlabel=L"time $t$ [s]",box=true)
figname="Helicopter_experimental_voltages.svg"
savefig(figpath*figname)

U_th = ConstantInterpolation(u_th,tm)
U_psi = ConstantInterpolation(u_psi,tm)
Theta = ConstantInterpolation(theta,tm)
Psi = ConstantInterpolation(psi,tm);

# Set-up for simulating helicopter
#
th0 = Theta(0.0)
psi0 = Psi(0.0)
w_th0 = 0.0
w_psi0 = 0.0
#
x0 = [th0,psi0,w_th0,w_psi0]
u_t = (th=U_th, psi=U_psi)
#
helicopter! = (dx,x,Ps,t) -> helicopter_ode(dx,x,Ps,t,P0,u_t)
#
tspan = (tm[1],tm[end])
#
prob = ODEProblem(helicopter!,x0,tspan,Ps)
#
sol = solve(prob,AutoTsit5(Rosenbrock23()),reltol=1e-6)
#
plot(sol,vars=(rad2deg,0,1),lw=LW1,lc=LC1,label=L"$\theta$ [deg]")
plot!(sol,vars=(rad2deg,0,2),lw=LW1,lc=LC2,label=L"$\psi$ [deg]")
plot!(xlim = tspan,xlabel=L"time $t$ [s]", ylabel="angles")
plot!(title="Angles: experimental inputs, nominal parameters",box=true)
figname="Helicopter_simulated_angles_experimental_inputs_nominal.svg"
savefig(figpath*figname)

function cost(p,Ps,pidx,x0,xidx,data,Nd,tspan)
    #
    xx0 = copy(x0)
    xx0[xidx] = p[xidx].*x0[xidx]
    Ps[pidx] .= p[length(xidx)+1:end]
    #
    prob = ODEProblem(helicopter!,xx0,tspan,Ps)
    sol = solve(prob, AutoTsit5(Rosenbrock23()), saveat=dt)
    #
    return norm(sol[1:2,:]-data,2)^2/Nd
    #
end
#
;

#
# (:rc, :h, :m, :Jbc, :Js, :g, :K_th_th, :K_th_psi, :K_psi_th, :K_psi_psi, :D_th, :D_psi)
Ps_keys = [:rc, :h, :m, :Jbc, :Js, :K_th_th,:K_th_psi,:K_psi_th,:K_psi_psi,:D_th, :D_psi]
pidx = [findall(Ps_keys[i] .== keys(P0))[1] for i in 1:length(Ps_keys)]
#
xidx = collect(3:4)
#
p = fill(1.0,length(xidx)+length(pidx))
#
p_lo = 0.2*p
p_hi = 2*p
p_lo_hi = collect(zip(p_lo,p_hi))
#
# Initial value of loss function
Ps = fill(1.0,length(P0))
x0 = [Theta(0),Psi(0),1.0,1.0]
#
loss = (p) -> cost(p,Ps,pidx,x0,xidx,data,Nd,tspan)
;

loss(p)

@time res = bboptimize(loss; SearchRange = p_lo_hi, NumDimensions = length(p));

best_fitness(res)

p_best = best_candidate(res)

loss(p_best)

x0_best = copy(x0)
x0_best[xidx] = p_best[xidx].*x0[xidx]
Ps_best=fill(1.0,length(P0))
Ps_best[pidx] .= p_best[length(xidx)+1:end]
#
prob_best = ODEProblem(helicopter!,x0_best,tspan,Ps_best)
#
sol_best = solve(prob_best,AutoTsit5(Rosenbrock23()),reltol=1e-6)
#
plot(sol_best,vars=(rad2deg,0,1),lw=LW1,lc=LC1,label=L"$\theta$ [deg]")
plot!(tm,theta*180/pi,lw=LW1,ls=LS2,lc=LC2,label=L"\theta^\mathrm{d}")
plot!(xlim = tspan,xlabel=L"time $t$ [s]", ylabel=L"$\theta$ [$^\circ$]")
plot!(title="Pitch angle: model (blue, solid) vs. data (red, dotted)",box=true)
figname="Helicopter_pitch-angle_model_fit.svg"
savefig(figpath*figname)

plot(sol_best,vars=(rad2deg,0,2),lw=LW1,lc=LC1,label=L"$\psi$ [deg]")
plot!(tm,psi*180/pi,lw=LW1,ls=LS2,lc=LC2,label=L"\psi^\mathrm{d}")
plot!(xlim = tspan,xlabel=L"time $t$ [s]", ylabel=L"$\psi$ [$^\circ$]")
plot!(title="Yaw angle: model (blue, solid) vs. data (red, dotted)",box=true)
figname="Helicopter_yaw-angle_model_fit.svg"
savefig(figpath*figname)

# Estimated initial states -- states 3 and 4 have been estimated
x0_best[xidx] = p_best[xidx].*x0[xidx]
show(x0_best)

# Estimated model parameters -- all parameters except gravity g have been assumed uncertain
Ps_best=fill(1.0,length(P0))
Ps_best[pidx] .= p_best[length(xidx)+1:end]
collect(zip(keys(P0),Ps_best.*values(P0)))

# Nominal parameters, taken from laboratory task description (and rounded somewhat)
P0

open("../optimization_results/globaloptimizationresults.jl", "w") do f
    write(f,"P0 = $P0 \n")
    write(f,"Ps_best = $Ps_best \n")
    write(f,"x0_best = $x0_best \n")
    write(f,"tspan = $tspan \n")
end


# %% Roessler trajectory -> PSOS -> Lorenz map
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("colorscheme.jl"))
using DynamicalSystems
import GLMakie

N = 10000.0

ds = Systems.roessler()
tr = trajectory(ds, N; Ttr = 100.0)

fig = GLMakie.Figure(resolution = (2100, 650), )
display(fig)

# Plot trajectory
trplot = fig[1,1] = GLMakie.Axis3(fig, scenekw = (camera = GLMakie.cam3d!, raw = false))
GLMakie.lines!(trplot, columns(tr)...; color = COLORS[3], linewidth = 2.0)
trplot.elevation = 0.4826990816987242
trplot.azimuth = 4.045530633326985
trplot.zticklabelsvisible=false
trplot.xticklabelsvisible=false
trplot.yticklabelsvisible=false

# Plot plane and section
o = GLMakie.Point3f0(-10, 0, 0)
w = GLMakie.Point3f0(25, 0, 25)
p = GLMakie.FRect3D(o, w)
a = GLMakie.RGBAf0(0,0,0,0)
c = GLMakie.to_color(COLORS[4])
img = GLMakie.Makie.ImagePattern([c a; a c]); # This throws an error if it shows, you can ignore that
GLMakie.mesh!(trplot, p; color = img)

psos = poincaresos(ds, (2, 0.0), N; Ttr = 100.0)
GLMakie.scatter!(trplot, columns(psos)...; color = COLORS[1], markersize = 5000)

# Plot section separately
psplot = GLMakie.Axis(fig[1, 2])
GLMakie.scatter!(psplot, psos[:, 1], psos[:, 3]; color = COLORS[1])
psplot.xlabel = "xₙ"
psplot.ylabel = "zₙ"

LS = 50
TS = 40

psplot.xticklabelsize = TS
psplot.yticklabelsize = TS
psplot.xlabelsize = LS
psplot.ylabelsize = LS

# Plot lorenz map
loplot = fig[1, 3] = GLMakie.Axis(fig)
GLMakie.scatter!(loplot, psos[1:end-1, 3], psos[2:end, 3]; color = COLORS[1])
loplot.xlabel = "zₙ"
loplot.ylabel = "zₙ₊₁"
loplot.xticklabelsize = TS
loplot.yticklabelsize = TS
loplot.xlabelsize = LS
loplot.ylabelsize = LS

trplot.zlabelsize = TS
trplot.xlabelsize = TS
trplot.ylabelsize = TS
trplot.xlabeloffset = 20
trplot.ylabeloffset = 20
trplot.zlabeloffset = 20

# GLMakie.Label(fig[1, 1, GLMakie.TopRight()], "b");

display(fig)

GLMakie.save(plotsdir("4", "roessler_map.png"), fig)