module RestrictProlong

using Base: tail

export restrict, restrict!, prolong

const DimsLike = Union{Dims,AbstractVector{Int}}

### restrict, for reducing the image size by 2-fold

"""
    Ar = restrict(A[, dims])

Perform two-fold reduction in size along the dimensions listed in
`dims`, or all coordinates if `dims` is omitted.  It anti-aliases A,
so is more accurate than a naive summation over
2×2×... blocks. `restrict` is normalized so as to approximately
preserve the mean value of `A`.

Thought of as an operator, `restrict` is equal to the transpose of `prolong`.
"""
restrict(A::AbstractArray, dims::DimsLike=1:ndims(A)) =  _restrict(A, dims...)

restrict(A::AbstractArray, dim::Integer) = _restrict(A, dim)

@inline _restrict(A, d1, d2, d...) = _restrict(_restrict(A, d1), d2, d...)
_restrict(A) = A

"""
    Ap = prolong(A, [sz::Tuple])

Perform two-fold expansion in size. If `sz` is omitted, all
dimensions are expanded. If `sz` is specified, it can be one of the
following:

- `sz[d] = size(A, d)`: no expansion is performed
- `sz[d] = 2*size(A, d) - 1`: expansion to an odd-size, adding (by interpolation) points on the half-grid between all points along dimension `d`
- `sz[d] = 2*size(A, d)`: expansion to an even-size, estimating the points on the 1/4 and 3/4 grid (including one point beyond either edge of the current span)

Any other choice results in an error. The default choice is to perform
odd-expansion. `prolong` is normalized so as to approximately preserve
the sum of `A`.

Thought of as an operator, `prolong` is equal to the transpose of `restrict`.
"""
prolong(A::AbstractArray) = prolong(A, map(n->2*n-1, _size(A)))

prolong(A::AbstractVector, sz::Integer) = prolong(A, sz, 1)

prolong(A::AbstractArray{T, N}, sz::Dims{N}) where {T, N} = _prolong(A, sz, ntuple(identity, Val(N)))

@inline _prolong(A, sz, dims) = _prolong(prolong(A, sz[1], dims[1]), tail(sz), tail(dims))
_prolong(A, ::Tuple{}, ::Tuple{}) = A

### Inner routines

function _restrict(A::AbstractArray{T, N}, dim::Integer) where {T, N}
    indsA = axes(A)
    sz = dim <= N ? length(indsA[dim]) : 1
    if sz <= 2
        return copy(A)
    end
    newinds = ntuple(i -> i==dim ? oftype(indsA[i], Base.OneTo(restrict_size(sz))) : indsA[i], Val(N))
    A1 = first(A)
    out = similar(A, typeof(A1/4+A1/2), newinds)
    restrict!(out, A, dim)
    out
end

# This properly anti-aliases. The only "oddity" is that edges tend
# towards zero under repeated restriction.

function restrict!(out, A, dim)
    checkdims(out, A, dim)
    indsA = axes(A)
    Rpre = CartesianIndices(indsA[1:dim-1])
    Rpost = CartesianIndices(indsA[dim+1:end])
    _restrict!(out, axes(out, dim), A, Rpre, indsA[dim], Rpost)
end

# Normalized so that out has roughly the same mean as A
@noinline function _restrict!(out::AbstractArray{T}, indout, A, Rpre::CartesianIndices, indA, Rpost::CartesianIndices) where T
    l = length(indA)
    fill!(out, zero(T))
    if isodd(l)
        half = convert(eltype(T), 0.5)
        quarter = convert(eltype(T), 0.25)
        rngA = first(indA):2:last(indA)-1
        rngout = first(indout):last(indout)-1
        for Ipost in Rpost
            for (i,j) in zip(rngA, rngout)
                @inbounds for Ipre in Rpre
                    tmp = quarter*A[Ipre, i+1, Ipost]
                    out[Ipre, j, Ipost]   += half*A[Ipre, i, Ipost] + tmp
                    out[Ipre, j+1, Ipost] += tmp
                end
            end
            @inbounds for Ipre in Rpre
                out[Ipre, last(indout), Ipost] += half*A[Ipre, last(indA), Ipost]
            end
        end
    else
        threeeighths = convert(eltype(T), 0.375)
        oneeighth = convert(eltype(T), 0.125)
        half = convert(eltype(T), 0.5)
        rngA = first(indA)+1:2:last(indA)-1
        rngout = first(indout):last(indout)-1
        for Ipost in Rpost
            @inbounds for Ipre in Rpre
                out[Ipre, first(indout), Ipost] = half*A[Ipre, first(indA), Ipost]
            end
            @inbounds for (i,j) in zip(rngA, rngout)
                for Ipre in Rpre
                    tmpA1, tmpA2 = A[Ipre, i, Ipost], A[Ipre, i+1, Ipost]
                    out[Ipre, j, Ipost]   += threeeighths*tmpA1 + oneeighth*tmpA2
                    out[Ipre, j+1, Ipost] += oneeighth*tmpA1 + threeeighths*tmpA2
                end
            end
            @inbounds for Ipre in Rpre
                out[Ipre, last(indout), Ipost] += half*A[Ipre, last(indA), Ipost]
            end
        end
    end
    out
end

"""
    prolong(A, sz::Integer, dim::Integer)

Perform expansion to size `sz` along dimension `dim`. `sz` must be one
of the valid choices for prolongation size.
"""
function prolong(A::AbstractArray{T, N}, sz::Integer, dim::Integer) where {T, N}
    if dim > N
        sz == 1 || throw(DimensionMismatch("cannot prolong $N-dimensional array to size $sz along dimension $dim (size must be 1)"))
        return copy(A)
    end
    indsA = axes(A)
    l = dim <= ndims(A) ? length(indsA[dim]) : 1
    (sz == l) | (sz == 2*l-1) | (sz == 2*l) || throw(DimensionMismatch("along dimension $dim, sz must be one of $l, $(2*l-1), or $(2*l), got $sz"))
    newinds = ntuple(i -> i==dim ? oftype(indsA[i], Base.OneTo(sz)) : indsA[i], Val(N))
    A1 = first(A)
    out = similar(A, typeof(A1/2+A1/2), newinds)
    Rpre = CartesianIndices(indsA[1:dim-1])
    Rpost = CartesianIndices(indsA[dim+1:end])
    _prolong!(out, newinds[dim], A, Rpre, indsA[dim], Rpost)
end

@noinline function _prolong!(out::AbstractArray{T}, indout, A, Rpre::CartesianIndices, indA, Rpost::CartesianIndices) where T
    l = length(indout)
    fill!(out, zero(T))
    if isodd(l)
        half = convert(eltype(T), 0.5)
        quarter = convert(eltype(T), 0.25)
        rngout = first(indout):2:last(indout)-1
        for Ipost in Rpost
            i = first(indA)
            for j in rngout
                @inbounds for Ipre in Rpre
                    tmpA1, tmpA2 = A[Ipre, i, Ipost], A[Ipre, i+1, Ipost]
                    out[Ipre, j, Ipost]   = half*tmpA1
                    out[Ipre, j+1, Ipost] = quarter*tmpA1+quarter*tmpA2
                end
                i += 1
            end
            @inbounds for Ipre in Rpre
                out[Ipre, last(indout), Ipost] = half*A[Ipre, last(indA), Ipost]
            end
        end
    else
        threeeighths = convert(eltype(T), 0.375)
        oneeighth = convert(eltype(T), 0.125)
        half = convert(eltype(T), 0.5)
        rngout = first(indout)+1:2:last(indout)-1
        rngA = first(indA):last(indA)-1
        for Ipost in Rpost
            @inbounds for Ipre in Rpre
                out[Ipre, first(indout), Ipost] = half*A[Ipre, first(indA), Ipost]
            end
            @inbounds for (i,j) in zip(rngA, rngout)
                for Ipre in Rpre
                    tmpA1, tmpA2 = A[Ipre, i, Ipost], A[Ipre, i+1, Ipost]
                    out[Ipre, j, Ipost]   = threeeighths*tmpA1 + oneeighth*tmpA2
                    out[Ipre, j+1, Ipost] = oneeighth*tmpA1 + threeeighths*tmpA2
                end
            end
            @inbounds for Ipre in Rpre
                out[Ipre, last(indout), Ipost] = half*A[Ipre, last(indA), Ipost]
            end
        end
    end
    out
end

restrict_size(len::Integer) = (len+1)>>1

function checkdims(out, A, dim)
    1 <= dim <= ndims(A) || throw(DimensionMismatch("dim must be between 1 and $(ndims(A)), got $dim"))
    ndims(out) == ndims(A) || throw(DimensionMismatch("A and out must have the same dimensions, got $(ndims(A)) and $(ndims(out))"))
    indsA = axes(A)
    indsout = axes(out)
    for i = 1:ndims(A)
        if i != dim
            indsout[i] == indsA[i] || throw(DimensionMismatch("A and out must have the same axces for all non-reduced dimensions, got $(indsA[i]) and $(indsout[i]) for dimension $i"))
        end
    end
    indAr = indsA[dim]
    indout = indsout[dim]
    length(indout) == restrict_size(length(indAr)) || throw(DimensionMismatch("A and out disagree along the restricted dimension, have axes $indAr and $indout but latter must be of length $(length(indout))"))
    nothing
end

_size(A::AbstractArray) = map(length, axes(A))

end # module


using Gadfly

set_default_plot_size(6inch, 6inch)

import Base: convert, show, +, -, /, *, isless, one, zero, isfinite

struct Percent
    value::Float64
end

# Functions necessary for plotting

+(a::Percent, b::Percent) = Percent(a.value + b.value)
-(a::Percent, b::Percent) = Percent(a.value - b.value)
-(a::Percent) = Percent(-a.value)
*(a::Percent, b::Float64) = Percent(a.value * b)
*(a::Float64, b::Percent) = Percent(a * b.value)

# Must return something that can be converted to Float64 with float64(a/b)
/(a::Percent, b::Percent) = a.value / b.value

isless(a::Percent, b::Percent) = isless(a.value, b.value)
one(::Type{Percent}) = Percent(0.01)
zero(::Type{Percent}) = Percent(0.0)
isfinite(a::Percent) = isfinite(a.value)
convert(::Type{Float64}, x::Percent) = x.value
show(io::IO, p::Percent) = print(io, round(100 * p.value, digits=4), "%")

y=[Percent(0.1), Percent(0.2), Percent(0.3)]
plot(x=collect(1:length(y)), y=y)


#!/bin/env julia

# Playing with running the pricer on GPU
#
# Will port into main library once I'm happy with it

module CUDA

export mc_pricer, mc_pricer_pathdep

using CuArrays # Unfortunately this gives a warning - see https://github.com/JuliaPackaging/Requires.jl/issues/65
CuArrays.allowscalar(false)

using Statistics: mean, std

present_value(v,r,t,T) = v*CuArrays.CUDAnative.exp(-r*(T-t))

function rand_price(seed=randn();payoff=(S,K,t,T,r)->S,S_0=100,t=0,T=1,K=S_0,r=0.02,σ=0.05) # Doesn't work on GPU - doesn't like function as argument
    S = B(seed;T=T-t,B0=S_0,r=r,d=0,σ=σ)

    present_value(payoff(S,K,t,T,r),r,t,T)
end

function mc_pricer(payoff;rng=CURAND.randn,S_0=100,t=0,T=1.0,r=0.02,σ=0.05,trials=10_000)
  a = rng(Float32,trials)
  #a = B.(rng(Float32,trials),0.9) # ;B0=S_0,T=T-t,r=0.02,σ=0.05) # CuArrays doesn't like any extra args?

  # TODO: stuff gets converted to f64 here - which IIRC is much slower on GPUs
  Bbake(x) = S_0*CuArrays.CUDAnative.exp(r*T-0.5*σ^2*T+σ*CuArrays.CUDAnative.pow(T-t,0.5)*x) # sqrt(T-t) doesn't work
  prices = present_value.(payoff.(Bbake.(a)),r,t,T)
  (mean = mean(prices), sem = std(prices)/sqrt(trials))
end

# TODO:
# - this gives different results to JustJoshing.mc_pricer_pathdep - build a failing test
# - try to unify the interface with the CPU based one (tricky as we need to "reach in" and specify dims=2?)
function mc_pricer_pathdep(payoff;S_0=100,t=0:0.1:1,T=1.0f0,r=0.02f0,σ=0.05f0,trials=100_000)
  # Much faster than CPU - now check why!
  a = CuArrays.CURAND.randn(Float32,trials)

  Bbake(p,t)::Float32 = begin
    S_0*CuArrays.CUDAnative.exp(
        r*T -
        0.5f0*σ^2*T +
        p*σ*CuArrays.CUDAnative.sqrt(T-t)
    )
  end
  # If prices is not a cuarray we get a big slowdown
  # Need to ensure that our function operates on array for max
  prices::CuArray{Float32,2} = payoff(Bbake.(a,t'))

  (mean = mean(prices), sem = std(prices)/sqrt(trials))
end

end


using StochasticDiffEq, Plots, DiffEqDevTools, DiffEqProblemLibrary
using ParameterizedFunctions, DiffEqMonteCarlo, Base.Test, OrdinaryDiffEq
using Plots; plotly()
using BenchmarkTools

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

## Lorenz

f = @ode_def_nohes LorenzSDE begin
  dx = σ*(y-x)
  dy = x*(ρ-z) - y
  dz = x*y - β*z
end σ ρ β

σ = (du,u,p,t) -> begin
  for i in 1:3
    du[i] = 10.0 #Additive
  end
end

p = (10.0,28.0,2.66)
prob_sde_lorenz = SDEProblem(f,σ,ones(3),(0.0,10.0),p)

monte_prob = MonteCarloProblem(prob_sde_lorenz)
N = 100; fails = zeros(5,4); times = zeros(5,4)
offset = 0
for i in 1:5
  sol1 =solve(monte_prob,SRA1();num_monte=N,abstol=10,reltol=1/2^(2),verbose=false)
  sol2 =solve(monte_prob,SRA3();num_monte=N,abstol=10,reltol=1/2^(offset+i),verbose=false)
  sol3 =solve(monte_prob,SOSRA();num_monte=N,abstol=10,reltol=1/2^(offset+i),verbose=false)
  sol4 =solve(monte_prob,SOSRA2();num_monte=N,abstol=10,reltol=1/2^(offset+i),verbose=false)
  fails[i,1] = N - sum([sol1[i].retcode == :Success for i in 1:N])
  fails[i,2] = N - sum([sol2[i].retcode == :Success for i in 1:N])
  fails[i,3] = N - sum([sol3[i].retcode == :Success for i in 1:N])
  fails[i,4] = N - sum([sol4[i].retcode == :Success for i in 1:N])
  times[i,1] = sol1.elapsedTime
  times[i,2] = sol2.elapsedTime
  times[i,3] = sol3.elapsedTime
  times[i,4] = sol4.elapsedTime
end


using JuMP
using Combinatorics: permutations

"""Stores and validates the values of a deviation for discrete probabilities `p`.

# Arguments
- `n::Int`: Length of discrete probabilities
- `p::Vector{Float64}`: Vector of discrete probabilities
- `d⁻::Vector{Float64}`: Vector of lower bounds
- `d⁺::Vector{Float64}`: Vector of uppper bounds
- `ϵ::Float64`: Radius of uncertainty

The arguments should satisfy the following conditions:

```julia
n ≥ 1
all(p .≥ 0)
isapprox(sum(p), 1)
length(p) == length(d⁻) == length(d⁺) == n
all(d⁻ .≥ -p)
all(d⁺ .≤ 1.0.-p)
0 ≤ ϵ ≤ 1
```

# Examples
Arguments

```julia
n = 2
p = [0.4, 0.6]
d⁻ = [-0.2, -0.2]
d⁺ = [0.2, 0.2]
ϵ = 0.1
```
"""
struct Deviation
    n::Int
    p::Vector{Float64}
    d⁻::Vector{Float64}
    d⁺::Vector{Float64}
    ϵ::Float64
    function Deviation(n::Int, p::Vector{Float64}, d⁻::Vector{Float64}, d⁺::Vector{Float64}, ϵ::Float64)
        n ≥ 1 || throw(DomainError(""))
        all(p .≥ 0) || throw(DomainError(""))
        isapprox(sum(p), 1) || throw(DomainError(""))
        length(p) == length(d⁻) == length(d⁺) == n || throw(DomainError(""))
        all(d⁻ .≥ -p) || throw(DomainError(""))
        all(d⁺ .≤ 1.0.-p) || throw(DomainError(""))
        0 ≤ ϵ ≤ 1 || throw(DomainError(""))
        new(n, p, d⁻, d⁺, ϵ)
    end
end

"""Constructs a deviation.

Examples
```julia
Deviation(p, d⁻, d⁺, ϵ)
```
"""
function Deviation(p::Vector{Float64}, d⁻::Vector{Float64}, d⁺::Vector{Float64}, ϵ::Float64)
    n = length(p)
    Deviation(n, p, d⁻, d⁺, ϵ)
end

"""Constructs a deviation.

```julia
Deviation(p, d⁻, d⁺)
```
"""
function Deviation(p::Vector{Float64}, d⁻::Vector{Float64}, d⁺::Vector{Float64})
    Deviation(p, d⁻, d⁺, 1.0)
end

"""Constructs a deviation.

```julia
Deviation(p, ϵ)
```
"""
function Deviation(p::Vector{Float64}, ϵ::Float64)
    Deviation(p, -p, 1.0.-p, ϵ)
end

"""Computes the optimal cross-assignment recursively."""
function cross_assignment(l::Int, h::Int, Σ::Float64, d::Vector{Float64}, d⁻::Vector{Float64}, d⁺::Vector{Float64}, ϵ::Float64)
    if (h - l ≤ 0) || (ϵ-Σ ≤ 0)
        return d
    end
    δ_l = d⁺[l] - d[l]
    δ_h = d[h] - d⁻[h]
    δ = min(δ_l, δ_h, ϵ-Σ)
    d[l] += δ
    d[h] -= δ
    if δ_l < δ_h
        return cross_assignment(l+1, h, Σ+δ, d, d⁻, d⁺, ϵ)
    elseif δ_l > δ_h
        return cross_assignment(l, h-1, Σ+δ, d, d⁻, d⁺, ϵ)
    else
        return cross_assignment(l+1, h-1, Σ+δ, d, d⁻, d⁺, ϵ)
    end
end

"""Compute cross-assignment."""
function cross_assignment(k::Int, d⁻::Vector{Float64}, d⁺::Vector{Float64}, ϵ::Float64)
    d = zeros(k)
    cross_assignment(1, k, 0.0, d, d⁻, d⁺, ϵ)
end

"""Computes the optimal cross assignment from deviation.

# Examples
```julia
dev = Deviation(...)
cross_assignment(dev)
```
"""
function cross_assignment(dev::Deviation)
    cross_assignment(dev.n, dev.d⁻, dev.d⁺, dev.ϵ)
end

"""Polyhedral uncertainty"""
function polyhedral_uncertainty(perm::Vector{Int}, dev::Deviation)
    q = copy(dev.p)
    q[perm] += cross_assignment(length(perm), dev.d⁻[perm], dev.d⁺[perm], dev.ϵ)
    return q
end

"""Polyhedral uncertainty set."""
function polyhedral_uncertainty_set(dev::Deviation)::Array{Vector{Float64}}
    p, ϵ = dev.p, dev.ϵ
    if iszero(ϵ)
        return [p]
    else
        i = [i for i in LinearIndices(p) if !iszero(p[i])]
        s = Set(polyhedral_uncertainty(perm, dev) for perm in permutations(i))
        return collect(s)
    end
end


function solve()
    x, n = [parse(Int, x) for x in split(readline())]
    p = Set(parse(Int, x) for x in split(readline()))

    d = 0
    while true
        for s in [-1, 1]
            a = x + d*s
            if a ∉ p
                return a
            end
        end
        d += 1
    end

end

println(solve())


# make your scripts automatically re-activate your project
cd(@__DIR__)
using Pkg; Pkg.activate("../..")

# load packages
using ml_for_pt
using DelimitedFiles
using LaTeXStrings
using Plots
using Random
ENV["GKSwstype"]="nul"

# set linear system size and corresponding data folder
L = 10
# L=60
data_folder = "../../data/ising/L="*string(L)*"/"

# set path to save folder
save_folder = "./results/L="*string(L)*"/"

# critical parameter value for the Ising model
p_crit = Float32(2/log(1+sqrt(2)))

# define parameter ranges
p_min = 0.05f0
p_max = 10.0f0
dp = 0.05f0
p_range = collect(p_min:dp:p_max)
p_range_LBC = collect(p_min-dp/2:dp:p_max+dp/2)

# import Monte Carlo data

# unique energies at p=p_1
energy = readdlm(data_folder*"E_T="*string(0)*".txt", eltype(dp))

# number of times each energy is drawn at T=T_1
num = readdlm(data_folder*"num_T="*string(0)*".txt", eltype(dp))

# iterate over entire parameter range
energies=[energy]
numbers=[num]
for index in 1:length(p_range)-1
  energyy = readdlm(data_folder*"E_T="*string(index)*".txt", eltype(dp))
  numm = readdlm(data_folder*"num_T="*string(index)*".txt", eltype(dp))
  push!(energies, energyy)
  push!(numbers, numm)
end
unique_energies = unique(vcat(energies...))
n_samples = sum(numbers[1])

# construct probability distribution from Monte Carlo data
# distr(x,p) gives probability to draw the sample x at parameter value p
# samples correspond to enumeration of unique energies with integers
distr,samples = MLP.distr_approx(energies, unique_energies, numbers, p_range)

# supervised learning using analytical expression
# returns optimal predictions, indicator, and loss
pred_opt_SL, indicator_opt_SL, loss_opt_SL = MLP.get_indicators_SL_analytical(samples, distr, p_range, dp, p_min, p_max)

# prediction-based method using analytical expression
# returns optimal predictions, indicator, and loss
pred_opt_PBM, indicator_opt_PBM, loss_opt_PBM = MLP.get_indicators_PBM_analytical(samples, distr, p_range, dp)

# learning by confusion using analytical expression
# returns optimal indicator and loss
indicator_opt_LBC, loss_opt_LBC = MLP.get_indicators_LBC_analytical(samples, distr, p_range, p_range_LBC)

# save output in save_folder
open(save_folder*"p_range_SL.txt", "w") do io
  writedlm(io, p_range)
end

open(save_folder*"p_range_PBM.txt", "w") do io
  writedlm(io, p_range)
end

open(save_folder*"p_range_LBC.txt", "w") do io
  writedlm(io, p_range_LBC)
end

open(save_folder*"pred_opt_SL.txt", "w") do io
  writedlm(io, pred_opt_SL)
end

open(save_folder*"indicator_opt_SL.txt", "w") do io
  writedlm(io, indicator_opt_SL)
end

open(save_folder*"loss_opt_SL.txt", "w") do io
  writedlm(io, [loss_opt_SL])
end

open(save_folder*"pred_opt_PBM.txt", "w") do io
  writedlm(io, pred_opt_PBM)
end

open(save_folder*"indicator_opt_PBM.txt", "w") do io
  writedlm(io, indicator_opt_PBM)
end

open(save_folder*"loss_opt_PBM.txt", "w") do io
  writedlm(io, [loss_opt_PBM])
end

open(save_folder*"indicator_opt_LBC.txt", "w") do io
  writedlm(io, indicator_opt_LBC)
end

open(save_folder*"loss_opt_LBC.txt", "w") do io
  writedlm(io, loss_opt_LBC)
end

# plot and save results in save_folder
plt = plot(p_range, pred_opt_SL, dpi=300, c="black", label=false)
vline!([p_crit], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$\hat{y}^{\mathrm{opt}}_{\mathrm{SL}}$")
savefig(save_folder*"pred_opt_SL.png")

plt = plot(p_range[2:end-1], indicator_opt_SL, dpi=300, c="black", label=false)
vline!([p_crit], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$I^{\mathrm{opt}}_{\mathrm{SL}}$")
savefig(save_folder*"indicator_opt_SL.png")

plt = plot(p_range, pred_opt_PBM, dpi=300, c="black", label=false)
vline!([p_crit], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$\hat{y}^{\mathrm{opt}}_{\mathrm{PBM}}$")
savefig(save_folder*"pred_opt_PBM.png")

plt = plot(p_range[2:end-1], indicator_opt_PBM, dpi=300, c="black", label=false)
vline!([p_crit], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$I^{\mathrm{opt}}_{\mathrm{PBM}}$")
savefig(save_folder*"indicator_opt_PBM.png")

plt = plot(p_range_LBC, indicator_opt_LBC, dpi=300, c="black", label=false)
vline!([p_crit], label=L"p_{c}",c="red")
xlabel!(L"$p$")
ylabel!(L"$I^{\mathrm{opt}}_{\mathrm{LBC}}$")
savefig(save_folder*"indicator_opt_LBC.png")

plt = plot(p_range_LBC, loss_opt_LBC, dpi=300, c="black",label=false)
vline!([p_crit], label=L"p_{c}",c="red")
xlabel!(L"$p$")
ylabel!(L"$\mathcal{L}^{\mathrm{opt}}_{\mathrm{LBC}}$")
savefig(save_folder*"loss_opt_LBC.png")


#
#
#

using Distributions
using QuadGK

betas_transition = [-4, -1.7, -1.2]#[-3.0, -2.5, -2.0]
sigma_eps = 0.5
alpha = 1.0
c = 0.5
delta = 0.9

function generate_simplex_3dims(n_per_dim::Integer=20)

    simplex = [[x, y, 1-x-y] for x in linspace(0,1,n_per_dim) for y in linspace(0,1,n_per_dim) if x+y <= 1.0]

    return hcat(simplex...)'

end

#= OLD interpV() method
using GridInterpolations

function interpV(simplex, V)

    step = simplex[2,2]
    grid = 0.0:step:1.0
    size_simplex = length(grid)

    augment_V = zeros(size_simplex, size_simplex)
    k = 1
    for i in 1:size_simplex
        for j in 1:size_simplex
            if j <= size_simplex+1-i
                augment_V[i,j] = V[k]
                k += 1
            end
        end
    end

    interpolate_V(x) = interpolate(SimplexGrid(grid, grid), flipdim(augment_V,1), [1-x[1], x[2]])

    return interpolate_V

end
=#

function tripolate(V, x, y)
    # Triangular (linear) interpolation:
    # 0<=|x|<=1, 0<=|y|<=1 and N parition points
    # of [0,1] for x, y
    N  = size(V, 1)-1
    xf = Int(floor(N*x) - floor(x))
    yf = Int(floor(N*y) - floor(y))
    x_ = (N*x - xf)
    y_ = (N*y - yf)
    if min(xf, yf) == 0 && max(xf, yf) <= 1
        x_ = 1-x_
        y_ = 1-y_
    end
    if xf + yf <= 1
        return V[xf+2,yf+1]*(1-x_) + V[xf+1,yf+2]*(1-y_) - V[xf+1,yf+1]*(1-x_-y_)
    else
        return V[xf+2,yf+1]*(x_) + V[xf+1,yf+2]*(y_) + V[xf+1,yf+1]*(1-x_-y_)
    end
end

function interpV(simplex, V)
	N_simplex = Int((sqrt(8*size(lambdas, 1)+1)-1)/2)
	augment_V = zeros(N_simplex, N_simplex)
    k = 1
    for i in 1:N_simplex
        for j in 1:(N_simplex+1-i)
            augment_V[i,j] = V[k]
            k += 1
        end
    end

	interpolate_V(x) = tripolate(augment_V, x[1], x[2])

	return interpolate_V
end

function dmd_transition_fs(new_state, action, old_state)
    return [pdf(Normal(alpha + beta*log.(action), sigma_eps), new_state) for beta in betas_transition]
end

function belief(new_state, action, old_state, lambda_weights)
    return dot(dmd_transition_fs(new_state, action, old_state), lambda_weights)
end

function update_lambdas(new_state, action, old_state, old_lambdas)
    den = dot(old_lambdas, dmd_transition_fs(new_state, action, old_state))
    return (dmd_transition_fs(new_state, action, old_state)/den).*old_lambdas
end

function period_return(price_grid, lambdas)
    constant_part = (price_grid - c) * exp(alpha + (sigma_eps^2)/2)
    summation     = exp.(log.(price_grid)*betas_transition') * lambdas

    #constant_part = (p-const.c) * np.e ** const.α * np.e ** ((const.σ_ɛ ** 2) / 2)
    #summation = np.dot(np.e**(betas_transition*np.log(p[:, np.newaxis])), lambdas)

    return (summation.*constant_part)

end

function myopic_price(lambdas)

    elasticity = dot(lambdas, betas_transition)
    @assert elasticity < -1.0

    return c/(1+(1/elasticity))

    #elasticity = np.dot(lambdas, betas_transition) #-2.2
    #assert elasticity < -1.0
    #return const.c / (1 + (1/elasticity))

end

function E0fV(Vguess, price_grid, lambda_weights)
    # Integrates Vguess * belief_f for each value of p. Integration over demand
    # Sum of points on demand and weights, multiplied by V and the belief function

    integrated_values = similar(price_grid)
    error = 0

    for i in 1:length(integrated_values)

        new_lambdas(x) = update_lambdas(x, price_grid[i], ~, lambda_weights)
        new_belief(x)  = belief(x, price_grid[i], ~, lambda_weights)

        #Vguess takes all lambdas except last (because of the simplex)
        integrand(x) = Vguess(new_lambdas(x))*new_belief(x)

        logd_min, logd_max = -6, 5 #D = (0.01, 10)
        integrated_values[i], error_tmp = quadgk(integrand, logd_min, logd_max, maxevals=2000)

        error += error_tmp

    end

    @assert error < 10.0^-4

    return integrated_values

end

function bellman_operator(Vguess, price_grid, lambda_simplex)

    policy = zeros(size(lambda_simplex,1))
    T_V    = zeros(size(lambda_simplex,1))

    # 1. Go over grid of state space
    # 2. Write objective (present return + delta*eOfV)
    # 3. Find optimal p on that objective
    # 4. Write optimal p and value function on that point in the grid

    for index in 1:size(lambda_simplex, 1)

        #println("> Doing ", index, " of ", size(lambda_simplex,1))

        lambda_weights = lambda_simplex[index, :]

        R      = period_return(price_grid, lambda_weights)
        E0fV_p = E0fV(Vguess, price_grid, lambda_weights)

        objective_vals = R + delta *  E0fV_p

        T_V[index], ind_max = findmax(objective_vals)
        policy[index]       = price_grid[ind_max]

    end

    # Interpolate
    interp_T_V = interpV(lambda_simplex, T_V)

    return interp_T_V, policy

end

function compute_fixed_point(V, price_grid, lambda_simplex; error_tol=1e-5, max_iter=50, verbose=true, skip=10)

    iterate = 1
    error = error_tol + 1

    policy = similar(price_grid)
    error  = Inf

    while iterate <= max_iter && error > error_tol
        if verbose && (mod(iterate, skip) == 0)
            tic()
        end

        new_V, policy = bellman_operator(V, price_grid, lambda_simplex)

        # Compute error over simplex:
        error = 0
        for i in 1:size(lambda_simplex,1)
            error += abs(V(lambda_simplex[i,:])-new_V(lambda_simplex[i,:]))
        end

        V = new_V

        if verbose && (mod(iterate, skip) == 0)
            println(@sprintf("Computed iterate %d with error %.4f", iterate, error))
            print(" !-- ")
            toc()
        end

        iterate += 1

    end

    return V, policy, error

end


using VPTrees, DataStructures
using Base.Test

function euclidean(coords::Matrix, i::Int, j::Int)
    @inbounds x = coords[1, i] - coords[1, j]
    @inbounds y = coords[2, i] - coords[2, j]
    @inbounds z = coords[3, i] - coords[3, j]

    return sqrt(x^2 + y^2 + z^2)
end

function euclidean_periodic(coords::Matrix, i::Int, j::Int)
    xbw, ybw, zbw = 1.0, 1.0, 1.0

    xhbw = xbw / 2
    yhbw = ybw / 2
    zhbw = zbw / 2

    @inbounds x = coords[1, i] - coords[1, j]
    @inbounds y = coords[2, i] - coords[2, j]
    @inbounds z = coords[3, i] - coords[3, j]

    if x < -xhbw x += xbw elseif x > xhbw x -= xbw end
    if y < -yhbw y += ybw elseif y > yhbw y -= ybw end
    if z < -zhbw z += zbw elseif z > zhbw z -= zbw end

    return sqrt(x^2 + y^2 + z^2)
end

function is_consistent{T}(metric::Function, node::Node{T}, point_index::Int, τ::T, isless)
    node.index == 0 && return true

    μ = metric(node.index, point_index)
    #println("$(point_index) and $(node.index): $μ <= $τ is $(μ <= τ) and isless = $isless")
    ((μ <= τ) ⊻ isless) && return false

    node.isleaf && return true

    is_consistent(metric, node.inside, point_index, τ, isless) || return false
    is_consistent(metric, node.outside, point_index, τ, isless) || return false

    return true
end

function is_consistent{T}(metric::Function, node::Node{T})
    node.isleaf && return true

    is_consistent(metric, node.inside, node.index, node.radius, true) || return false
    is_consistent(metric, node.outside, node.index, node.radius, false) || return false
    #println()
    is_consistent(metric, node.inside) || return false
    is_consistent(metric, node.outside) || return false

    return true
end

is_consistent(tree::VPTree) = is_consistent(tree.metric, tree.root)

create_metric(distance::Function, points) = (i::Int, j::Int) -> distance(points, i, j)

function test_knn(distance::Function, num_points::Int, num_iters::Int)
    for s in 1:num_iters
        srand(s)
        points = rand(3, num_points)
        point_indices = collect(1:num_points)
        metric = create_metric(distance, points)

        num_neighbors = 5

        tree = VPTree(metric, point_indices)
        neighbors = NeighborList{Float64}(num_neighbors)

        for i in 1:num_points
            knn!(neighbors, tree, i)
            indices_test = indices(neighbors)

            pq = PriorityQueue(Int, Float64, Base.Order.Reverse)
            for j in 1:num_points
                i == j && continue

                r = metric(i, j)

                if length(pq) < num_neighbors
                    enqueue!(pq, j, r)
                elseif r < last(peek(pq))
                    dequeue!(pq)
                    enqueue!(pq, j, r)
                end
            end

            indices_actual = keys(pq)

            isempty(symdiff(indices_test, indices_actual)) || return false
        end
    end

    return true
end

function test_rangesearch(distance::Function, num_points::Int, num_iters::Int)
    for s in 1:num_iters
        srand(s)
        points = rand(3, num_points)
        point_indices = collect(1:num_points)
        metric = create_metric(distance, points)

        tree = VPTree(metric, point_indices)
        neighbors = NeighborList{Float64}(num_points)

        r = rand()
        for i in 1:num_points
            rangesearch!(neighbors, tree, i, r)
            indices_test = indices(neighbors)

            indices_actual = filter(j -> metric(i, j) <= r, 1:num_points)

            isempty(symdiff(indices_test, indices_actual)) || return false
        end
    end

    return true
end

function test_integrity(distance::Function, num_points::Int, num_iters::Int)
    for s in 1:num_iters
        srand(s)
        points = rand(3, num_points)
        point_indices = collect(1:num_points)
        metric = create_metric(distance, points)

        tree = VPTree(metric, point_indices)
        is_consistent(tree) || return false
    end

    return true
end

@testset "VP Tree tests" begin
    const num_points = 1000
    const num_iters = 5

    @testset "VP Tree integrity test" begin
        @test test_integrity(euclidean, num_points, num_iters)
        @test test_integrity(euclidean_periodic, num_points, num_iters)
    end

    @testset "k-nearest neighbors test" begin
        @test test_knn(euclidean, num_points, num_iters)
        @test test_knn(euclidean_periodic, num_points, num_iters)
    end

    @testset "Range search test" begin
        @test test_rangesearch(euclidean, num_points, num_iters)
        @test test_rangesearch(euclidean_periodic, num_points, num_iters)
    end
end


using AQFED, Test
import AQFED.PLN: EtoreGobetPLNPricer, LeFlochLehmanPLNPricer, priceEuropean
import AQFED.TermStructure: CapitalizedDividend, Dividend, futureValue
import AQFED.Basket: DeelstraBasketPricer, DeelstraLBBasketPricer

@testset "EtoreGobetSingle" begin
    spot = 100.0
    vol = 0.3
    discountRate = 0.0
    divAmount = 7.0
    tte = 1.0
    ttd = 0.9 * tte
    ttp = tte
    isCall = true
    discountFactor = exp(-discountRate * ttp)
    dividends = Vector{CapitalizedDividend}(undef, 1)
    #amount::T        exDate::Float64        payDate::Float64        isProportional::bool        isKnown::bool
    dividends[1] =
        CapitalizedDividend(Dividend(divAmount, ttd, ttd, false, false), exp((tte - ttd) * discountRate))
    println("S Model Price Error")
    eg2 = EtoreGobetPLNPricer(2)
    eg3 = EtoreGobetPLNPricer(3)

    ll2 = LeFlochLehmanPLNPricer(2)
    ll3 = LeFlochLehmanPLNPricer(3)
    d = DeelstraBasketPricer(1, 3)
    dlb = DeelstraLBBasketPricer(1, 3)
    refHHL = [43.24845580, 9.07480013, 1.06252880, 0.10473887]
    refEg2 = [43.24846889, 9.07479039, 1.06253441, 0.10474035]
    refEg3 = [43.24845582, 9.07480026, 1.06252875, 0.10473885]
    for (i, strike) in enumerate(LinRange(50.0, 200.0, 4))
        price = priceEuropean(
            eg2,
            isCall,
            strike,
            spot * exp(tte * discountRate),
            vol * vol * tte,
            tte,
            discountFactor,
            dividends,
        )
        println(strike, " EG-2 ", price, " ", price - refHHL[i])
        @test isapprox(refEg2[i], price, atol = 1e-8)
        price = priceEuropean(
            eg3,
            isCall,
            strike,
            spot * exp(tte * discountRate),
            vol * vol * tte,
            tte,
            discountFactor,
            dividends,
        )
        println(strike, " EG-3 ", price, " ", price - refHHL[i])
        @test isapprox(refEg3[i], price, atol = 1e-8)
        price = priceEuropean(
            ll2,
            isCall,
            strike,
            spot * exp(tte * discountRate),
            vol * vol * tte,
            tte,
            discountFactor,
            dividends,
        )
        println(strike, " LL-2 ", price, " ", price - refHHL[i])
        price = priceEuropean(
            ll3,
            isCall,
            strike,
            spot * exp(tte * discountRate),
            vol * vol * tte,
            tte,
            discountFactor,
            dividends,
        )
        println(strike, " LL-3 ", price, " ", price - refHHL[i])
        price = AQFED.Basket.priceEuropean(
            d,
            isCall,
            strike,
            spot * exp(tte * discountRate),
            vol * vol * tte,
            tte,
            discountFactor,
            dividends,
        )
        println(strike, " Deelstra ", price, " ", price - refHHL[i])
        price = AQFED.Basket.priceEuropean(
            dlb,
            isCall,
            strike,
            spot * exp(tte * discountRate),
            vol * vol * tte,
            tte,
            discountFactor,
            dividends,
        )
        println(strike, " Deelstra-LB ", price, " ", price - refHHL[i])

    end

end

@testset "MultipleVellekoop" begin
    spot = 100.0
    σ = 0.25
    discountRate = 0.06
    tte = 7.0
    isCall = true
    ttd = 0.9
    dividends = [CapitalizedDividend(Dividend(6.0, ttd, ttd, false, false), exp((tte - ttd) * discountRate)),
    CapitalizedDividend(Dividend(6.5, ttd + 1, ttd + 1, false, false), exp((tte - ttd - 1) * discountRate)),
    CapitalizedDividend(Dividend(7.0, ttd + 2, ttd + 2, false, false), exp((tte - ttd - 2) * discountRate)),
    CapitalizedDividend(Dividend(7.5, ttd + 3, ttd + 3, false, false), exp((tte - ttd - 3) * discountRate)),
    CapitalizedDividend(Dividend(8.0, ttd + 4, ttd + 4, false, false), exp((tte - ttd - 4) * discountRate)),
    CapitalizedDividend(Dividend(8.0, ttd + 5, ttd + 5, false, false), exp((tte - ttd - 5) * discountRate)),
    CapitalizedDividend(Dividend(8.0, ttd + 6, ttd + 6, false, false), exp((tte - ttd - 6) * discountRate))]
    rawForward = spot * exp(discountRate * tte)
    df = exp(-discountRate * tte)
    ll3 = LeFlochLehmanPLNPricer(3)
    d = DeelstraBasketPricer(3, 3)
    dlb = DeelstraLBBasketPricer(3, 3)
    refPrices = [
        34.19664044,
        30.49396147,
        27.21393776,
        24.31297892,
        21.74885502,
        19.48228598,
        17.47762235,
        15.70300783,
        14.13025908,
        12.73460326,
        11.49435473,
    ]
    f = spot * exp(discountRate * tte) - sum(futureValue(cd) for cd in dividends)

    for (i, strike) in enumerate(LinRange(50.0, 150.0, 11))
        refVol = Black.impliedVolatility(isCall, refPrices[i], f, strike, tte, df)
        price = priceEuropean(ll3, isCall, strike, rawForward, σ^2 * tte, tte, df, dividends)
        vol = Black.impliedVolatility(isCall, price, f, strike, tte, df)
        println(strike, " LL-3 ", price, " ", price / refPrices[i] - 1, " ", vol - refVol)
        price = AQFED.Basket.priceEuropean(d, isCall, strike, rawForward, σ^2 * tte, tte, df, dividends)
        vol = Black.impliedVolatility(isCall, price, f, strike, tte, df)
        println(strike, " Deelstra ", price, " ", price / refPrices[i] - 1, " ", vol - refVol)
        price = AQFED.Basket.priceEuropean(dlb, isCall, strike, rawForward, σ^2 * tte, tte, df, dividends)
        vol = Black.impliedVolatility(isCall, price, f, strike, tte, df)
        println(strike, " Deelstra-LB ", price, " ", price / refPrices[i] - 1, " ", vol - refVol)

    end
end

@testset "MultipleGocsei" begin
    spot = 100.0
    σ = 0.25
    discountRate = 0.03
    divAmount = 2.0
    tte = 10.0
    isCall = true
    ttd0 = 1.0 / 365
    dividends = Vector{CapitalizedDividend}(undef, 20)
    for i = 1:20
        ttd = ttd0 + (i - 1) / 2
        dividends[i] =
            CapitalizedDividend(Dividend(divAmount, ttd, ttd, false, false), exp((tte - ttd) * discountRate))
    end
    rawForward = spot * exp(discountRate * tte)
    df = exp(-discountRate * tte)
    ll3 = LeFlochLehmanPLNPricer(3)
    d = DeelstraBasketPricer(1, 3) #1,3 minimizes relative error in price, 3,3 minimizes volatility error
    dlb = DeelstraLBBasketPricer(1, 3)
    refPrices = [37.25772649, 28.59516740, 22.28192756, 17.61042951, 14.09731856, 11.41453970, 9.33695198]
    f = spot * exp(discountRate * tte) - sum(futureValue(cd) for cd in dividends)

    for (i, strike) in enumerate(LinRange(50.0, 200.0, 7))
        refVol = Black.impliedVolatility(isCall, refPrices[i], f, strike, tte, df)
        price = priceEuropean(ll3, isCall, strike, rawForward, σ^2 * tte, tte, df, dividends)
        vol = Black.impliedVolatility(isCall, price, f, strike, tte, df)
        println(strike, " LL-3 ", price, " ", price / refPrices[i] - 1, " ", vol - refVol)
        price = AQFED.Basket.priceEuropean(d, isCall, strike, rawForward, σ^2 * tte, tte, df, dividends)
        vol = Black.impliedVolatility(isCall, price, f, strike, tte, df)
        println(strike, " Deelstra ", price, " ", price / refPrices[i] - 1, " ", vol - refVol)
        price = AQFED.Basket.priceEuropean(dlb, isCall, strike, rawForward, σ^2 * tte, tte, df, dividends)
        vol = Black.impliedVolatility(isCall, price, f, strike, tte, df)
        println(strike, " Deelstra-LB ", price, " ", price / refPrices[i] - 1, " ", vol - refVol)

    end
end


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ 9f4d1108-1fed-4314-ae4e-560ec7a9be9c
using PlutoTest

# ╔═╡ 76d27f60-b0e8-11eb-3fd3-a7b73d0d4cfe
md"""
# [Day 14: Reindeer Olympics](https://adventofcode.com/2015/day/14)

This year is the Reindeer Olympics! Reindeer can fly at high speeds, but must rest occasionally to recover their energy. Santa would like to know which of his reindeer is fastest, and so he has them race.

Reindeer can only either be flying (always at their top speed) or resting (not moving at all), and always spend whole seconds in either state.

For example, suppose you have the following Reindeer:

-   Comet can fly 14 km/s for 10 seconds, but then must rest for 127 seconds.
-   Dancer can fly 16 km/s for 11 seconds, but then must rest for 162 seconds.

After one second, Comet has gone 14 km, while Dancer has gone 16 km. After ten seconds, Comet has gone 140 km, while Dancer has gone 160 km. On the eleventh second, Comet begins resting (staying at 140 km), and Dancer continues on for a total distance of 176 km. On the 12th second, both reindeer are resting. They continue to rest until the 138th second, when Comet flies for another ten seconds. On the 174th second, Dancer flies for another 11 seconds.

In this example, after the 1000th second, both reindeer are resting, and Comet is in the lead at `1120` km (poor Dancer has only gotten `1056` km by that point). So, in this situation, Comet would win (if the race ended at 1000 seconds).

Given the descriptions of each reindeer (in your puzzle input), after exactly `2503` seconds, **what distance has the winning reindeer traveled?**
"""

# ╔═╡ 7b3cffcf-5c04-4d13-9714-8dad9e4956f3
puzzle_input = "Vixen can fly 19 km/s for 7 seconds, but then must rest for 124 seconds.
Rudolph can fly 3 km/s for 15 seconds, but then must rest for 28 seconds.
Donner can fly 19 km/s for 9 seconds, but then must rest for 164 seconds.
Blitzen can fly 19 km/s for 9 seconds, but then must rest for 158 seconds.
Comet can fly 13 km/s for 7 seconds, but then must rest for 82 seconds.
Cupid can fly 25 km/s for 6 seconds, but then must rest for 145 seconds.
Dasher can fly 14 km/s for 3 seconds, but then must rest for 38 seconds.
Dancer can fly 3 km/s for 16 seconds, but then must rest for 37 seconds.
Prancer can fly 25 km/s for 6 seconds, but then must rest for 143 seconds."

# ╔═╡ 4c30acd3-078f-40a6-8e22-c90704abef38
function parseline(line)
	m = match(r"(\w+) can fly (\d+) km/s for (\d+) seconds, but then must rest for (\d+) seconds.", line)
	(km=parse(Int, m[2]), fly_seconds=parse(Int, m[3]), rest_seconds=parse(Int, m[4]))
end

# ╔═╡ 6259b41c-e73c-46b6-bf76-2f42c6428003
example_players = map(parseline, ["Comet can fly 14 km/s for 10 seconds, but then must rest for 127 seconds.", "Dancer can fly 16 km/s for 11 seconds, but then must rest for 162 seconds."])

# ╔═╡ 21d4c9cf-2502-4522-90fc-5136e5bd13bd
function distance(player, time)
	d = 0
	#while time > 0
	#	stepTime = min(fly_seconds, time)
	#	d += stepTime * km
	#	time -= stepTime + rest_seconds
	#end

	d = time ÷ (player.fly_seconds + player.rest_seconds) * (player.fly_seconds * player.km)
	d += min(time % (player.fly_seconds + player.rest_seconds), player.fly_seconds) * player.km
	d
end

# ╔═╡ bec1d683-76e6-4b80-968c-4179944073cc
maxdistance(players, time) = maximum(map(player->distance(player, time), players))

# ╔═╡ cac38821-4cef-4d15-a86b-7878a37392da
@assert maxdistance(example_players, 1000) == 1120

# ╔═╡ 315d6498-c7d6-44f6-876f-3cbc599d7393
puzzle_players = map(parseline, split(puzzle_input, "\n"))

# ╔═╡ 598ca1b1-acfc-449d-be51-6ae63efaeaa7
part1 = maxdistance(puzzle_players, 2503)

# ╔═╡ 40ae0576-75ae-4d63-b163-4fbefe0e0060
@test part1 == 2660

# ╔═╡ 17210bdf-21cd-4f48-8e86-b7c568ee3e2d
md"Your puzzle answer was `2660`."

# ╔═╡ 0975d5a4-675e-42b8-8ab1-4b31dc2de102
md"""
# Part Two

Seeing how reindeer move in bursts, Santa decides he's not pleased with the old scoring system.

Instead, at the end of each second, he awards one point to the reindeer currently in the lead. (If there are multiple reindeer tied for the lead, they each get one point.) He keeps the traditional 2503 second time limit, of course, as doing otherwise would be entirely ridiculous.

Given the example reindeer from above, after the first second, Dancer is in the lead and gets one point. He stays in the lead until several seconds into Comet's second burst: after the 140th second, Comet pulls into the lead and gets his first point. Of course, since Dancer had been in the lead for the 139 seconds before that, he has accumulated 139 points by the 140th second.

After the 1000th second, Dancer has accumulated `689` points, while poor Comet, our old champion, only has `312`. So, with the new scoring system, Dancer would win (if the race ended at 1000 seconds).

Again given the descriptions of each reindeer (in your puzzle input), after exactly `2503` seconds, **how many points does the winning reindeer have?**
"""

# ╔═╡ 9cd8509d-38c3-49d0-bc6d-84c862bd37e0
function maxscore(players, time)
	scores = Dict()

	for idx in 1:length(players)
		scores[idx] = 0
	end

	for t in 1:time
		distances = Dict()
		for idx in 1:length(players)
			distances[idx] = distance(players[idx], t)
		end
		maxdistance = maximum(values(distances))
		winners = filter(idx-> distances[idx] == maxdistance, keys(distances))
		for idx in winners
			scores[idx] += 1
		end
	end
	maximum(values(scores))

end

# ╔═╡ 14d5831a-f848-426d-a014-2a1e24ee1d68
@assert maxscore(example_players, 1000) == 689

# ╔═╡ 0dec0d65-0cc9-4d7c-bdf0-d2bbc3f6c956
part2 = maxscore(puzzle_players, 2503)

# ╔═╡ 09982bbe-252b-40e2-b818-c1f2f575f458
@test part2 == 1256

# ╔═╡ fa92058e-35cf-43b8-8552-71dd16e1c52e
md"Your puzzle answer was `1256`."

# ╔═╡ Cell order:
# ╠═9f4d1108-1fed-4314-ae4e-560ec7a9be9c
# ╟─76d27f60-b0e8-11eb-3fd3-a7b73d0d4cfe
# ╟─7b3cffcf-5c04-4d13-9714-8dad9e4956f3
# ╠═4c30acd3-078f-40a6-8e22-c90704abef38
# ╠═6259b41c-e73c-46b6-bf76-2f42c6428003
# ╠═21d4c9cf-2502-4522-90fc-5136e5bd13bd
# ╠═bec1d683-76e6-4b80-968c-4179944073cc
# ╠═cac38821-4cef-4d15-a86b-7878a37392da
# ╠═315d6498-c7d6-44f6-876f-3cbc599d7393
# ╠═598ca1b1-acfc-449d-be51-6ae63efaeaa7
# ╠═40ae0576-75ae-4d63-b163-4fbefe0e0060
# ╟─17210bdf-21cd-4f48-8e86-b7c568ee3e2d
# ╟─0975d5a4-675e-42b8-8ab1-4b31dc2de102
# ╠═9cd8509d-38c3-49d0-bc6d-84c862bd37e0
# ╠═14d5831a-f848-426d-a014-2a1e24ee1d68
# ╠═0dec0d65-0cc9-4d7c-bdf0-d2bbc3f6c956
# ╠═09982bbe-252b-40e2-b818-c1f2f575f458
# ╟─fa92058e-35cf-43b8-8552-71dd16e1c52e


using Turing

const n_samples = 10^4

#Sample from transformed distributions, check moments.
within_4_stdev(dist::Distribution, samples) = begin
    s = sqrt(var(dist))
    m = mean(dist)
    n = length(samples)
    stdevs = 4
    @assert 0.95*s < sqrt(var(samples)) < 1.05*s
    @assert m-stdevs*s/sqrt(n) < mean(samples) < m+stdevs*s/sqrt(n)
end

macro test_dist(Dist)
    dist = symbol(lowercase(string(Dist.args[1])))
    args = Dist.args[2:end]
    quote
        within_4_stdev($Dist, $dist(State(), $(args...); dims=($n_samples,)))
    end
    #dist = deepcopy(Dist)
    #dist.args[1] = symbol(lowercase(string(dist.args[1])))
    #insert!(dist.args, 2, :(State()))
    #push!(dist.args, :(($n_samples,)))
    #:(within_4_stdev($Dist, $dist))
end

@test_dist Normal(2, 3)
@test_dist LogNormal(1., 0.5)
#@test_dist Chisq(2)
#MvNormal
m = [1., 2.]
v = [5. -1; -1 3.]
samples = [mvnormal(State(), m, v) for i =1:n_samples]
println((m, mean(samples)))
println((v, cov(hcat(samples...)')))

@test_dist Exponential(3.)
@test_dist Bernoulli(0.7)
#@test_dist Uniform()
@test_dist Laplace(1., 3.)
@test_dist Categorical([0.1, 0.3, 0.6])
@test_dist Gamma(1., 3.)
@test_dist Beta(1., 3.)
#
#Dirichlet
println(); println(); println();
alpha = [1., 1.5, 0.5]
dist = Dirichlet(alpha)
samples = [dirichlet(State(), alpha) for i =1:n_samples]
println((mean(dist), mean(samples)))
println((cov(dist), cov(hcat(samples...)')))

#Multidimensional sampling fallbacks work
laplace(State(), ones(3), 1.)
laplace(State(), ones(3), ones(3))
gamma(State(), ones(3), 1.)
gamma(State(), ones(3), ones(3))


function metropolisBase(curr::Real, log_prob::Function, stepSD::Real)
  cand = curr + randn() * stepSD
  logU = log(rand())
  logP = log_prob(cand) - log_prob(curr)
  accept = logP > logU
  draw = accept ? cand : curr
  return (draw, accept)
end

function metropolis(curr::Real, log_prob::Function, stepSD::Real)
  return metropolisBase(curr, log_prob, stepSD)[1]
end

function metropolis(curr::AbstractVector{<:Real}, log_prob::Function,
                    propCov::AbstractMatrix{<:Real})
  cand = rand(MvNormal(curr, propCov))
  p = log_prob(cand) - log_prob(curr)
  return p > log(rand()) ? cand : curr
end

# TODO: Check.
# See section 2 of http://probability.ca/jeff/ftpdir/adaptex.pdf
function metropolisAdaptive(curr::AbstractVector{<:Real}, log_prob::Function,
                            tuner::MvTuner)
  # Update summary stats.
  update!(curr, tuner)

  # Proposal covariance.
  propCov = if tuner.iter <= 2*tuner.d || tuner.beta > rand()
    0.01 * eye(tuner.d) / tuner.d
  else
    5.6644 * tuner.current_cov / tuner.d
  end

  # Metropolis step
  return metropolis(curr, log_prob, propCov)
end

"""
Adaptive metropolis (within Gibbs). See section 3 of the paper below:
  http://probability.ca/jeff/ftpdir/adaptex.pdf

Another useful website:
  https://m-clark.github.io/docs/ld_mcmc/index_onepage.html
"""
function metropolisAdaptive(curr::Real, log_prob::Function,
                            tuner::TuningParam;
                            update::Function=update_tuning_param_default)
  draw, accept = metropolisBase(curr, log_prob, tuner.value)
  update(tuner, accept)
  return draw
end

"""
Log absolute jacobian term to be added when taking the log of a postiviely
supported random variable (x), so as to produce a random variable with
support on the real line. When added to the target log density, the resulting
expression can be used in a metropolis step conveniently.
"""
function logabsjacobian_logx(real_x::Real)
  return real_x
end

"""
Log absolute jacobian term to be added when taking the logit of a
lower-and-upper-bounded random variable (x), so as to produce a random variable
with support on the real line. When added to the target log density, the
resulting expression can be used in a metropolis step conveniently.

# Arguments
- `real_x::Float64`: the unconstrained value of the positively-supported
                     parameter x (i.e. log x)
- `a::Float64`: lower bound for unconstrained parameter (default: 0)
- `b::Float64`: upper bound for unconstrained parameter (default: 1)
"""
function logabsjacobian_logitx(real_x::Real)
  return logpdf(Logistic(), real_x)
end

function metLogAdaptive(curr::Real, log_prob::Function, tuner::TuningParam;
                        update::Function=update_tuning_param_default)

  function log_prob_plus_logabsjacobian(log_x::Real)
    x = exp(log_x)
    return log_prob(x) + logabsjacobian_logx(log_x)
  end

  log_x = metropolisAdaptive(log(curr), log_prob_plus_logabsjacobian, tuner,
                             update=update)

  return exp(log_x)
end

function metLogitAdaptive(curr::Real, log_prob::Function,
                          tuner::TuningParam;
                          update::Function=update_tuning_param_default)

  function log_prob_plus_logabsjacobian(logit_x::Real)
    x = logistic(logit_x)
    return log_prob(x) + logabsjacobian_logitx(logit_x)
  end

  logit_x = metropolisAdaptive(logit(curr),
                               log_prob_plus_logabsjacobian,
                               tuner, update=update)

  return logistic(logit_x)
end


using SimulacoesTemplate
using Test

@testset "Distance function" begin

   # Testing rpbc
   x = [ 0, 0 ]
   y = [ 12, 12 ]
   side = 10
   @test SimulacoesTemplate.rpbc(x,y,side) ≈ sqrt(8)
   y = [ 17, 17 ]
   @test SimulacoesTemplate.rpbc(x,y,side) ≈ sqrt(18)

end

@testset "Potential Energy" begin

  data = Data()

  @test SimulacoesTemplate.upair(π,data) ≈ -0.0003250451775148634

  p = initial_point(data.N, data.side, 0.9)
  @test utotal(p,data) ≈ -7472.164849146468

end


# fast robust geometric predicates
# Converted into Julia from C
# Original functions written by
#  Jonathan Richard Shewchuk
#  School of Computer Science
#  Carnegie Mellon University
#  5000 Forbes Avenue
#  Pittsburgh, Pennsylvania  15213-3891
#  jrs@cs.cmu.edu

using GeometryTypes

"""
original comment:
/* Which of the following two methods of finding the absolute values is      */
/*   fastest is compiler-dependent.  A few compilers can inline and optimize */
/*   the fabs() call; but most will incur the overhead of a function call,   */
/*   which is disastrously slow.  A faster way on IEEE machines might be to  */
/*   mask the appropriate bit, but that's difficult to do in C.              */
This function definition attempts to do the same thing as a C-Define
"""
@inline absolute(a) = ((a) >= 0.0 ? (a) : -(a))
# TODO: should also try the built in absolute
# @inline absolute(a) = abs(a)

# Skipping tails, condensing to one function (?)
#@inline function fast_two_sum_tail(a::Float64, b::Float64, x::Float64)::Tuple{Float64,Float64}
#    bvirt::Float64 = x - a
#    x, b - bvirt
#end

@inline function fast_two_sum(a::Float64, b::Float64)::Tuple{Float64,Float64}
    x::Float64 = a + b
    bvirt::Float64 = x - a
    x, b - bvirt
end

#@inline function fast_two_diff_tail(a::Float64, b::Float64, x::Float64)

const resulterrbound = Ref{Float64}()
const ccwerrboundA = Ref{Float64}()
const ccwerrboundB = Ref{Float64}()
const ccwerrboundC = Ref{Float64}()
const o3derrboundA = Ref{Float64}()
const o3derrboundB = Ref{Float64}()
const o3derrboundC = Ref{Float64}()
const iccerrboundA = Ref{Float64}()
const iccerrboundB = Ref{Float64}()
const iccerrboundC = Ref{Float64}()
const isperrboundA = Ref{Float64}()
const isperrboundB = Ref{Float64}()
const isperrboundC = Ref{Float64}()
"""
/*****************************************************************************/
/*                                                                           */
/*  exactinit()   Initialize the variables used for exact arithmetic.        */
/*                                                                           */
/*  `epsilon' is the largest power of two such that 1.0 + epsilon = 1.0 in   */
/*  floating-point arithmetic.  `epsilon' bounds the relative roundoff       */
/*  error.  It is used for floating-point error analysis.                    */
/*                                                                           */
/*  `splitter' is used to split floating-point numbers into two half-        */
/*  length significands for exact multiplication.                            */
/*                                                                           */
/*  I imagine that a highly optimizing compiler might be too smart for its   */
/*  own good, and somehow cause this routine to fail, if it pretends that    */
/*  floating-point arithmetic is too much like real arithmetic.              */
/*                                                                           */
/*  Don't change this routine unless you fully understand it.                */
/*                                                                           */
/*****************************************************************************/
"""
function exactinit()::Nothing
    every_other::Bool = false
    epsilon::Float64 = 1.0
    splitter::Float64 = 1.0
    check::Float64 = 1.0
    #/* Repeatedly divide `epsilon' by two until it is too small to add to    */
    #/*   one without causing roundoff.  (Also check if the sum is equal to   */
    #/*   the previous sum, for machines that round up instead of using exact */
    #/*   rounding.  Not that this library will work on such machines anyway. */
    lastcheck = check;
    epsilon *= 0.5
    splitter *= 2.0
    check = 1.0 + epsilon
    while ((check != 1.0) && (check != lastcheck))
        lastcheck = check;
        epsilon *= 0.5
        if (every_other)
            splitter *= 2.0
        end
        every_other = !every_other
        check = 1.0 + epsilon
    end
    splitter += 1.0

    #/* Error bounds for orientation and incircle tests. */
    global resulterrbound[] = (3.0 + 8.0 * epsilon) * epsilon
    global ccwerrboundA[] = (3.0 + 16.0 * epsilon) * epsilon
    global ccwerrboundB[] = (2.0 + 12.0 * epsilon) * epsilon
    global ccwerrboundC[] = (9.0 + 64.0 * epsilon) * epsilon * epsilon
    global o3derrboundA[] = (7.0 + 56.0 * epsilon) * epsilon
    global o3derrboundB[] = (3.0 + 28.0 * epsilon) * epsilon
    global o3derrboundC[] = (26.0 + 288.0 * epsilon) * epsilon * epsilon
    global iccerrboundA[] = (10.0 + 96.0 * epsilon) * epsilon
    global iccerrboundB[] = (4.0 + 48.0 * epsilon) * epsilon
    global iccerrboundC[] = (44.0 + 576.0 * epsilon) * epsilon * epsilon
    global isperrboundA[] = (16.0 + 224.0 * epsilon) * epsilon
    global isperrboundB[] = (5.0 + 72.0 * epsilon) * epsilon
    global isperrboundC[] = (71.0 + 1408.0 * epsilon) * epsilon * epsilon
    nothing
end

"""
/*               Return a positive value if the point pd lies below the      */
/*               plane passing through pa, pb, and pc; "below" is defined so */
/*               that pa, pb, and pc appear in counterclockwise order when   */
/*               viewed from above the plane.  Returns a negative value if   */
/*               pd lies above the plane.  Returns zero if the points are    */
/*               coplanar.  The result is also a rough approximation of six  */
/*               times the signed volume of the tetrahedron defined by the   */
/*               four points.                                                */
"""
function orient3d(pa::Point3f0, pb::Point3f0, pc::Point3f0, pd::Point3f0)::Float64
    ad::Point3f0 = pa .- pd
    bd::Point3f0 = pb .- pd
    cd::Point3f0 = pc .- pd

    bdxcdy::Float64 = bd[1] * cd[2]
    cdxbdy::Float64 = cd[1] * bd[2]
    cdxady::Float64 = cd[1] * ad[2]
    adxcdy::Float64 = ad[1] * cd[2]
    adxbdy::Float64 = ad[1] * bd[2]
    bdxady::Float64 = bd[1] * ad[2]

    det::Float64 = adz * (bdxcdy - cdxbdy) +
                   bdz * (cdxady - adxcdy) +
                   cdz * (adxbdy - bdxady)
    permanent = (absolute(bdxcdy) + absolute(cdxbdy)) * absolute(adz) +
                (absolute(cdxady) + absolute(adxcdy)) * absolute(bdz) +
                (absolute(adxbdy) + absolute(bdxady)) * absolute(cdz)
    # TODO: define predetermined error bounds
    errbound = o3derrboundA * permanent
    if (det > errbound) || (-det > errbound)
        return det
    else
        println("WARNING: possible errors in floating point arithmetic!!!")
        return det
    end
    # TODO: implement robust predicates
    #return orient3dadapt(pa, pb, pc, pd, permanent)
end

"""
/*               Return a positive value if the point pe lies inside the     */
/*               sphere passing through pa, pb, pc, and pd; a negative value */
/*               if it lies outside; and zero if the five points are         */
/*               cospherical.  The points pa, pb, pc, and pd must be ordered */
/*               so that they have a positive orientation (as defined by     */
/*               orient3d()), or the sign of the result will be reversed.    */
"""
function insphere(pa::Point3f0, pb::Point3f0, pc::Point3f0, pd::Point3f0, pe::Point3f0)::Float64
    ae::Point3f0 = pa .- pe
    be::Point3f0 = pb .- pe
    ce::Point3f0 = pc .- pe
    de::Point3f0 = pd .- pe

    aexbey = ae[1] * be[2]
    bexaey = be[1] * ae[2]
    ab = aexbey - bexaey
    bexcey = be[1] * ce[2]
    cexbey = ce[1] * be[2]
    bc = bexcey - cexbey
    cexdey = ce[1] * de[2]
    dexcey = de[1] * ce[2]
    cd = cexdey - dexcey
    dexaey = de[1] * ae[2]
    aexdey = ae[1] * de[2]
    da = dexaey - aexdey

    aexcey = ae[1] * ce[2]
    cexaey = ce[1] * ae[2]
    ac = aexcey - cexaey
    bexdey = be[1] * de[2]
    dexbey = de[1] * be[2]
    bd = bexdey - dexbey

    abc = ae[3] * bc - be[3] * ac + ce[3] * ab
    bcd = be[3] * cd - ce[3] * bd + de[3] * bc
    cda = ce[3] * da + de[3] * ac + ae[3] * cd
    dab = de[3] * ab + ae[3] * bd + be[3] * da

    alift = sum(ae .* ae)
    blift = sum(be .* be)
    clift = sum(ce .* ce)
    dlift = sum(de .* de)

    det = (dlift * abc - clift * dab) + (blift * cda - alift * bcd)
    permanent = ((absolute(cexdey) + absolute(dexcey)) * absolute(bez) +
                 (absolute(dexbey) + absolute(bexdey)) * absolute(cez) +
                 (absolute(bexcey) + absolute(cexbey)) * absolute(dez)) * alift +
                ((absolute(dexaey) + absolute(aexdey)) * absolute(cez) +
                 (absolute(aexcey) + absolute(cexaey)) * absolute(dez) +
                 (absolute(cexdey) + absolute(dexcey)) * absolute(aez)) * blift +
                ((absolute(aexbey) + absolute(bexaey)) * absolute(dez) +
                 (absolute(bexdey) + absolute(dexbey)) * absolute(aez) +
                 (absolute(dexaey) + absolute(aexdey)) * absolute(bez)) * clift +
                ((absolute(bexcey) + absolute(cexbey)) * absolute(aez) +
                 (absolute(cexaey) + absolute(aexcey)) * absolute(bez) +
                 (absolute(aexbey) + absolute(bexaey)) * absolute(cez)) * dlift
    errbound = isperrboundA * permanent;
    if (det > errbound) || (-det > errbound)
        return det
    else
        println("WARNING: possible error in floating point arithmetic!!!")
        return det
    end
    # TODO: implement adaptive predicates
end


# Copyright 2017 Marcelo Forets
# See accompanying license file.

module Zonotopes

#===============================================================================
Basic types
===============================================================================#

"""
    Zonotope(c, gi)

A zonotope with given center and generators.

INPUT:

- c:  center
- gi: matrix whose columns are the generators

EXAMPLES:

Two-dimensional zonotope with floating point vertices:

    julia> Zonotope([0.9; 0], 0.1*eye(2))

    Zonotopes.Zonotope{Float64}([0.9,0.0],[0.1 0.0; 0.0 0.1])
"""
type Zonotope{T}
    c :: Vector{T}
    gi :: Matrix{T}
end
Zonotope() = Zonotope(Vector(), Matrix())

#===============================================================================
Methods
===============================================================================#

"
    α × Z

Scale a zonotope by a numerical factor.

EXAMPLES::

We scale by a factor of sqrt(2)::

    julia>  Z = Zonotope([0.9; 0], 0.1*eye(2))
    Zonotopes.Zonotope{Float64}([0.9,0.0],[0.1 0.0; 0.0 0.1])
    julia> sqrt(2) * Z
    Zonotopes.Zonotope{Float64}([1.27279,0.0],[0.141421 0.0; 0.0 0.141421])
"
function Base.:*(α::Number, Z::Zonotope)

    return Zonotope(α * Z.c, α * Z.gi)
end

"
    M × Z

Apply a linear map to a zonotope.

EXAMPLES::

Applying a matrix to a two-dimensional zonotope::

    julia> Z = Zonotope([0.9; 0], 0.1*eye(2))
    Zonotopes.Zonotope{Float64}([0.387,0.603],[0.043 0.0051; 0.067 0.0724])
    julia> M = [0.43 0.051; 0.67 0.724]
    2×2 Array{Float64,2}:
     0.43  0.051
     0.67  0.724
    julia> M*Z
    Zonotopes.Zonotope{Float64}([0.387,0.603],[0.043 0.0051; 0.067 0.0724])
"
function Base.:*(M::Matrix, Z::Zonotope)

    return Zonotope(M * Z.c, M * Z.gi)
end

"""
    Z1 + Z2

Minkowski addition of two zonotopes.
"""
function Base.:+(Z1::Zonotope, Z2::Zonotope)

    return Zonotope(Z1.c + Z2.c, [Z1.gi Z2.gi])
end

"""
    Z1 * Z2

Cartesian product of two zonotopes.
"""
function Base.:*(Z1::Zonotope, Z2::Zonotope)

    return Zonotope([Z1.c ; Z2.c], cat([1, 2], Z1.gi, Z2.gi))
end

#===============================================================================
Special Zonotopes
===============================================================================#

"""
    Box(center, radius)

A box (or hypercube in n-dimensions) with given center and radius.

INPUT:

- center : vector, the center of the box
- radius : number, the radius of the box

EXAMPLES:

Three-dimensional box centered at the origin:

    julia> Box([0., 0, 0], 1.1)
    Zonotopes.Zonotope{Float64}([0.0,0.0,0.0],[1.1 0.0 0.0; 0.0 1.1 0.0; 0.0 0.0 1.1])
"""
function Box(center::Vector, radius::Number)

    # dimension
    n = length(center)

    return Zonotope(center, radius * eye(n))
end

#===============================================================================
Input/Output
===============================================================================#

# - Plotting.
# - Saving and Loading zonotopes.
# - Interface with MATLAB's MPT (Multi-Parametric Toolbox).

#===============================================================================
Exports
===============================================================================#

# Types
export Zonotope

# Special zonotopes
export Box

end # module


export rotate
export getrotation
export applyconfig

rotate(g::Quaternion, q::Quaternion) = g * q
rotate(M::Array{<:Complex,2}, q::ComplexPlane) = ComplexPlane(M * vec(q))
rotate(p::ℝ³, q::Quaternion) = ℝ³(vec(q * Quaternion([0; vec(p)]) * conj(q))[2:4])
rotate(p::Array{ℝ³}, q::Quaternion) = map(x -> rotate(x, q), p)
getrotation(i::ℝ³, n::ℝ³) = begin
    if isapprox(i, n)
        return Quaternion(0, normalize(i))
    end
    u = normalize(cross(i, n))
    θ = acos(dot(normalize(i), normalize(n))) / 2
    Quaternion(θ, u)
end

"""
    applyconfig(p, q)

Apply a rigid body transformation in 3-space with the given array of points `p` and
configuration `q`.
"""
function applyconfig(p::Array{ℝ³}, q::Biquaternion)
    map(x -> x + gettranslation(q), rotate(p, getrotation(q)))
end


function FluidResistor(; name, R = 1.0)
    # 2 pin component
    @named oneport = OnePort()
    @unpack p, q = oneport
    ps = @parameters R=R
    eqs = [
        p ~ q * R
    ]
    extend(ODESystem(eqs, t, [], ps; name), oneport)
end

function VesselCompliance(; name, C = 1.0)
    @named oneport = OnePort()
    @unpack p, q = oneport
    ps = @parameters C=C
    D = Differential(t)
    eqs = [
        D(p) ~ q / C
    ]
    extend(ODESystem(eqs, t, [], ps; name), oneport)
end

function FluidInertia(; name, I = 1.0)
    @named oneport = OnePort()
    @unpack p, q = oneport
    ps = @parameters I=I
    D = Differential(t)
    eqs = [
        D(q) ~ p / I
    ]
    extend(ODESystem(eqs, t, [], ps; name), oneport)
end


module VortexLattice

using LinearAlgebra
using StaticArrays
using Rotations

## Panel geometry and math tools
#==========================================================================================#

import ..AeroMDAO: Panel3D, panel_area, panel_coords, midpoint, panel_normal, transform, p1, p2, p3, p4, average_chord, average_width, dynamic_pressure, aerodynamic_coefficients, force_coefficient, moment_coefficient, rate_coefficient, weighted_vector, reshape_array, cartesian_to_freestream, freestream_to_cartesian

## Horseshoe setup
#==========================================================================================#

include("horseshoes.jl")
include("finite_core.jl")

# export Horseshoe, horseshoe_line, horseshoe_point, bound_leg, bound_leg_center, bound_leg_vector, r1, r2, points, make_horseshoes

## Reference frames
#==========================================================================================#

include("reference_frames.jl")

# export body_to_stability_axes, stability_to_body_axes, body_to_wind_axes, wind_to_body_axes, reflect_xz

## Influence matrix and solution of system
#==========================================================================================#

include("influences.jl")

# export solve_system

"""
    make_horseshoes(horseshoe_panels)

Get the bound legs and collocation points of the horseshoes defined by horseshoe `Panel3D`s.
"""
make_horseshoes(horseshoe_panels) =	@. horseshoe_line(horseshoe_panels), horseshoe_point(horseshoe_panels)

"""
    solve_system(horseshoes, collocation_points, normals, U, Ω)

Evaluate and return the vortex strengths ``\\Gamma``s given `Horseshoes`, their associated normal vectors (not necessarily the same as the panels' normals), the speed ``U`` and rotation vector ``\\Omega``.
"""
function solve_system(horseshoes, normals, U, Ω)
    V = map(r -> U + Ω × r, collocation_point.(horseshoes))
    AIC  = influence_matrix(horseshoes, collocation_point.(horseshoes), normals, -normalize(U), false)
    boco = boundary_condition(V, normals)
    Γs 	 = AIC \ boco
end

## Force evaluations
#==========================================================================================#

include("nearfield.jl")

# export nearfield_dynamics, nearfield_drag

include("farfield.jl")

# export farfield_dynamics

## Post-processing
#==========================================================================================#

# export case_dynamics, evaluate_coefficients

"""
    case_dynamics(Γ_comp, hs_comp, Γs, horseshoes, U, α, β, Ω, ρ, r_ref)

Compute the nearfield and farfield forces over a given component with associated vortex strengths `Γ_comp` and horseshoes `hs_comp`, using the vortex strengths ``\\Gamma``s and `horseshoes` of the entire configuration, the speed ``U``, angles of attack ``\\alpha`` and sideslip ``\\beta``, the rotation vector ``\\Omega``, the freestream density ``\\rho`` and the reference location for moments ``r_\\text{ref}``.
"""
function case_dynamics(Γ_comp, hs_comp, Γs, horseshoes, U, α, β, Ω, ρ, r_ref)
    # Compute near-field dynamics
    surface_forces, surface_moments = nearfield_dynamics(Γ_comp, hs_comp, Γs, horseshoes, U, Ω, ρ, r_ref)

    # Compute farfield dynamics
    trefftz_force = trefftz_forces(Γ_comp, hs_comp, norm(U), α, β, ρ)

    surface_forces, surface_moments, trefftz_force
end

# In case of only one component
case_dynamics(Γs, horseshoes, U, α, β, Ω, ρ, r_ref) = case_dynamics(Γs, horseshoes, Γs, horseshoes, U, α, β, Ω, ρ, r_ref)

function evaluate_coefficients(forces, moments, trefftz_force, U, α, β, Ω, ρ, S, c, b)
    V = norm(U)
    q = dynamic_pressure(ρ, V)

    # Computing summed coefficients
    force, moment = sum(forces), sum(moments)

    # Transform near-field dynamics to wind axes
    trans_force	 = body_to_wind_axes(force, α, β)
    trans_force  = [ nearfield_drag(force, U); trans_force[2:end] ]
    trans_moment = body_to_wind_axes(stability_flip(moment), α, β)

    # Compute coefficients
    nearfield_coeffs = aerodynamic_coefficients(trans_force, trans_moment, Ω, V, S, b, c, ρ)
    farfield_coeffs	 = force_coefficient(trefftz_force, q, S)

    # Non-dimensional panel coefficients
    CFs = force_coefficient.(forces, q, S)
    CMs = moment_coefficient.(moments, q, S, b, c)

    nearfield_coeffs, farfield_coeffs, CFs, CMs
end

# Streamlines
include("streamlines.jl")

# export streamlines

# EXPERIMENTAL: SYSTEM
include("system.jl")

# export VLMState, VLMSystem, compute_influence_matrix!, compute_boundary_condition!, compute_collocation_points!, compute_horseshoes!, compute_normals!, solve_system!, compute_nearfield_forces!, compute_farfield_forces!, solve_case!, set_horseshoe_panels!, set_camber_panels!, compute_wake_properties!, normals, horseshoes, collocation_points, AIC, RHS, circulations, surface_force_coefficients, surface_moment_coefficients, total_force_coefficients, farfield_force_coefficients

end


import Pkg; Pkg.activate("../repo/TuringBnpBenchmarks")

using Turing
using Distributions
using DelimitedFiles
import Random

idx, Y = let
  data = readdlm("patients.csv", ',')
  Int.(data[:, 1]), data[:, 2:end]
end

# FAM without missing data.
@model fam(y, idx, K, L0, L1, phi=0) = begin
  I = length(unique(idx))
  Nsum = length(idx)
  J = size(y, 2)

  alpha ~ Gamma(1, 1)
  v ~ filldist(Beta(alpha/K, 1), K)
  Z ~ arraydist([Bernoulli(v[k]) for j in 1:J, k in 1:K])
  lambda = tzeros(Int, Nsum)

  sigma ~ filldist(LogNormal(0, 1), I)
  delta0 ~ filldist(LogNormal(0, 0.1), L0)
  delta1 ~ filldist(LogNormal(0, 0.1), L1)
  mu0 = -cumsum(delta0)
  mu1 = cumsum(delta1)

  W ~ filldist(Dirichlet(K, 1/K), I)

  eta0 = Array{Float64, 3}(undef, I, J, L0)
  eta1 = Array{Float64, 3}(undef, I, J, L1)
  for i in 1:I
    for j in 1:J
      eta0[i, j, :] ~ Dirichlet(L0, 1)
      eta1[i, j, :] ~ Dirichlet(L1, 1)
    end
  end

  for j in 1:J
    for n in 1:Nsum
      i = idx[n]

      lambda[n] ~ Categorical(W[:, i])
      z = Z[j, lambda[n]]
      Lz = (z == 0) ? L0 : L1
      muz = (z == 0) ? mu0 : mu1
      etaz = (z == 0) ? eta0 : eta1

      y[n, j] ~ UnivariateGMM(muz, fill(sigma[i], Lz), Categorical(etaz[i, j, :]))
    end
  end
end

Random.seed!(0)
@time chain = begin
    K = 25
    L0 = 5
    L1 = 5
    nburn = 10
    nsamps = 10
    iterations = nburn + nsamps
    sample(fam(Y, idx, K, L0, L1), Gibbs(PG(100)), iterations)
end;

nothing

#=
group(chain, :Z).value.data
=#


immutable AirspeedType{str} end
kcas = AirspeedType{:kcas}()
keas = AirspeedType{:keas}()
ktas = AirspeedType{:ktas}()
mach = AirspeedType{:mach}()

type Airspeed{T}
    val::Float64
    Airspeed(v) = typeof(T)<:AirspeedType ? new{T}(v) : error("Invalid Airspeed Type")
end

show{T}(io::IO,::AirspeedType{T}) = print(io,T)
show{T}(io::IO,a::Airspeed{T}) = print(io,a.val,T)
*{T}(x::Real,t::AirspeedType{T}) = Airspeed{t}(x)

immutable Env
    atm::Atm
    calibrated_airspeed::@quantity(Float64,m/s)
    equivalent_airspeed::@quantity(Float64,m/s)
    true_airspeed::@quantity(Float64,m/s)
    mach_number::Float64
    q::@quantity(Float64,N/m^2)
end

function calc_calibrated(mach_number::Float64,pressure_ratio::Float64)
     sqrt(5*((((0.2*mach_number^2+1)^3.5 - 1)*pressure_ratio+1)^(1/3.5) - 1))*SeaLevelSpeedOfSound
end

function calc_mach(calibrated_airspeed::@quantity(Float64,m/s),pressure_ratio::Float64)
    sqrt(5*((((0.2*(calibrated_airspeed/SeaLevelSpeedOfSound)^2+1)^3.5 - 1)/pressure_ratio+1)^(1/3.5) - 1))
end

function Env(airspeed::Airspeed{kcas},altitude::@quantity(Float64,m),delta_temperature::@quantity(Float64,K) = 0.0K)
    atm = Atm(altitude,delta_temperature)
    calibrated_airspeed = convert(SIQuantity,airspeed.val*knots)
    mach_number = calc_mach(calibrated_airspeed,atm.pressure_ratio)
    equivalent_airspeed = atm.speed_of_sound*mach_number
    true_airspeed = equivalent_airspeed/sqrt(atm.density_ratio)
    q = 0.5*atm.density*true_airspeed^2
    Env(atm,calibrated_airspeed,equivalent_airspeed,true_airspeed,mach_number,q)
end

function Env(airspeed::Airspeed{mach},altitude::@quantity(Float64,m),delta_temperature::@quantity(Float64,K) = 0.0K)
    atm = Atm(altitude,delta_temperature)
    mach_number = airspeed.val
    calibrated_airspeed = calc_calibrated(mach_number,atm.pressure_ratio)
    equivalent_airspeed = atm.speed_of_sound*mach_number
    true_airspeed = equivalent_airspeed/sqrt(atm.density_ratio)
    q = 0.5*atm.density*true_airspeed^2
    Env(atm,calibrated_airspeed,equivalent_airspeed,true_airspeed,mach_number,q)
end

function Env(airspeed::Airspeed{keas},altitude::@quantity(Float64,m),delta_temperature::@quantity(Float64,K) = 0.0K)
    atm = Atm(altitude,delta_temperature)
    equivalent_Airspeed = convert(SIQuantity,airspeed.val*knots)
    mach_number = equivalent_airspeed/atm.speed_of_sound
    calibrated_airspeed = calc_calibrated(mach_number,atm.pressure_ratio)
    true_airspeed = equivalent_airspeed/sqrt(atm.density_ratio)
    q = 0.5*atm.density*true_airspeed^2
    Env(atm,calibrated_airspeed,equivalent_airspeed,true_airspeed,mach_number,q)
end

# function Env{N1<:Real,N2<:Real,U1<:LengthUnit,U2<:TemperatureUnit}(airspeed::Airspeed,altitude::Quantity{N1,U1},delta_temperature::Quantity{N2,U2} = 0.0K)
#     Env(airspeed,convert(SIQuantity{Float64},altitude),convert(SIQuantity{Float64},delta_temperature))
# end


import Robotlib: ad, adi
function calibLPOE(xin,Tn0in,Ta,q;maxiter=10, λ=1.0)

    xi      = deepcopy(xin)
    Tn0     = deepcopy(Tn0in)
    n       = size(xi,2)-1
    N       = size(q,1)
    y       = zeros(6N)
    A       = zeros(6N,6*(n+1))
    xini    = deepcopy(xin)
    Tn0cand = deepcopy(Tn0)
    v       = zeros(3) # only revolute joints
    et      = zeros(maxiter+1)
    er      = zeros(maxiter+1)
    et[1],er[1]   = evalError(xi,Tn0,Ta,q)
    println("Error: ",et[1], " λ: ", λ)
    @assert size(Ta,3) == N

    for iter = 1:maxiter # do a few iterations of the calibration

        for i = 1:N
            Tfull = fkineLPOE(Tn0,xi,q[i,:]')
            # populate the matrices of the linear estimation problem
            y[(i-1)*6+1:6i] = twistcoords(logT(Ta[:,:,i]*trinv(Tfull)))
            A[(i-1)*6+1:6*i,1:6] = ad(Tn0[:,:,1])
            adi = ad(Tn0[:,:,1]*expξ(xi[:,1],q[i,1]))
            for j = 2:n
                A[(i-1)*6+1:6*i,(j-1)*6+1:6*j] = adi*ad(Tn0[:,:,j])
                adi = adi*ad(Tn0[:,:,j]*expξ(xi[:,j],q[i,j]))
            end
            A[(i-1)*6+1:6*i,n*6+1:6*(n+1)] = adi*ad(Tn0[:,:,n+1])
        end
        x = (A'A + λ*I)\A'y
        # Update the candidate nominal parameters
        for j = 1:n+1
            ξ = x[(j-1)*6+1:6j]
            Tn0cand[:,:,j] = Tn0[:,:,j]*expξ2(ξ,1)
        end

        et[iter+1],er[iter+1] = evalError(xi,Tn0cand,Ta,q)
        println("Error: ",round(et[iter+1],5), " λ: ", λ, " Norm dx: ", round(norm(x),5))

        if norm(x) < 1e-16
            return Tn0,xi, et, er
        end
        if et[iter+1] > et[iter]
            et[iter+1] = et[iter]
            er[iter+1] = er[iter]
            λ *= 10
        else
            λ/=10
            xi = deepcopy(xini)
            Tn0 = deepcopy(Tn0cand)
        end

    end

    return Tn0, xi, et, er

end

"""
`calibLPOEdual(xi,Tn0,q;maxiter=10, λ=1.0)`
Performs dual arm calibration given joint twists xi and nominal transformations Tn0
Returns calibrated nominal transformations
"""
function calibLPOEdual(xin,Tn0in,q;maxiter=10, λ=1.0)

    xi      = deepcopy(xin)
    Tn0     = deepcopy(Tn0in)
    n       = size(xi,2)-1
    N       = size(q,1)
    y       = zeros(6N)
    A       = zeros(6N,6*(n+1),2)
    xini    = deepcopy(xin)
    Tn0cand = deepcopy(Tn0)
    et      = zeros(maxiter+1)
    er      = zeros(maxiter+1)
    et[1],er[1]   = evalErrorDual(xi,Tn0,q)
    println("Error: ",et[1], " λ: ", λ)

    for iter = 1:maxiter # do a few iterations of the calibration
        for i = 1:N
            T1 = fkineLPOE(Tn0[:,:,:,1],xi[:,:,1],q[i,:,1]')
            T2 = fkineLPOE(Tn0[:,:,:,2],xi[:,:,2],q[i,:,2]') # Use the other arm as measurement
            y[(i-1)*6+1:6i] = twistcoords(logT(T2*trinv(T1)))
            for arm = 1:2
                A[(i-1)*6+1:6*i,1:6,arm] = ad(Tn0[:,:,1,arm])
                adi = ad(Tn0[:,:,1,arm]*expξ(xi[:,1,arm],q[i,1,arm]))
                for j = 2:n
                    A[(i-1)*6+1:6*i,(j-1)*6+1:6*j,arm] = adi*ad(Tn0[:,:,j,arm])
                    adi = adi*ad(Tn0[:,:,j,arm]*expξ(xi[:,j,arm],q[i,j,arm]))
                end
                A[(i-1)*6+1:6*i,n*6+1:6*(n+1),arm] = adi*ad(Tn0[:,:,n+1,arm])

            end
        end
        x = zeros(3)
        for arm = 1:2
            x = (A[:,:,arm]'A[:,:,arm] + λ*I)\A[:,:,arm]'y
            # Update the candidate nominal parameters
            s = arm == 1 ? 1 : -1 # Determine the sign of the correction
            for j = 1:n+1
                ξ = 0.5*s*x[(j-1)*6+1:6j] # Only apply half the update to each arm, verified to be essential
                Tn0cand[:,:,j,arm] = Tn0[:,:,j,arm]*expξ2(ξ,1)
            end

        end
        et[iter+1],er[iter+1] = evalErrorDual(xi,Tn0cand,q)
        println("Error: ",round(et[iter+1],5), " λ: ", λ, " Norm dx: ", round(norm(x),5))

        if norm(x) < 1e-14
            return Tn0,xi, et, er
        end
        if et[iter+1] > et[iter]
            et[iter+1] = et[iter]
            er[iter+1] = er[iter]
            λ *= 10
        else
            λ/=10
            xi = deepcopy(xini)
            Tn0 = deepcopy(Tn0cand)
        end

    end

    return Tn0, xi, et, er

end

function calibPOE(Xin,Ta,q;maxiter=50, λ = 10000.0)

    xin     = copy(Xin)
    n       = size(xin,2)-2
    N       = size(q,1)
    y       = zeros(6N)
    A       = zeros(6N,6*(n+1))
    xini    = similar(xin)
    et      = zeros(maxiter+1)
    er      = zeros(maxiter+1)
    et[1],er[1]   = evalErrorPOE(xin,Ta,q)
    println("Error: ",et[1], " λ: ", λ)
    @assert size(Ta,3) == N

    for iter = 1:maxiter # do a few iterations of the calibration

        for i = 1:N
            Tfull = fkinePOE(xin,q[i,:]')
            # populate the matrices of the linear estimation problem
            y[(i-1)*6+1:6i] = twistcoords(logT(Ta[:,:,i]*trinv(Tfull)))
            A[(i-1)*6+1:6*i,:] = Ai(q[i,:]',xin)

        end

        x = (A'A + λ*I)\A'y
        # Update the candidate nominal parameters
        for j = 1:n+1
            ξ = x[(j-1)*6+1:6j]
            xini[:,j] = xin[:,j] + ξ
            #xini[:,j] = conformize(xini[:,j])
        end

        et[iter+1],er[iter+1] = evalErrorPOE(xini, Ta,q)
        println("Error: ",round(et[iter+1],5), " λ: ", λ, " Norm dx: ", round(norm(x),5))

        if norm(x) < 1e-10
            return xin, et, er
        end
        if et[iter+1] > et[iter]
            et[iter+1] = et[iter]
            λ *= 10
        else
            λ/=10
            xin = deepcopy(xini)
        end

    end

    return xin, et, er

end

function calibPOE_offsets_from_points(Xin,Q;maxiter=50, λ = 10000.0)
    xin     = copy(Xin)
    n       = size(xin,2)-2
    Ndatasets = size(Q,1)
    Npoints = sum(map(x-> size(x,1),Q))

    y       = zeros(6Npoints)
    A       = zeros(6Npoints,n-1)
    xini    = similar(xin)
    et      = zeros(maxiter+1)
    er      = zeros(maxiter+1)
    et[1]   = evalErrorPOE_offsets_from_points(xin,Q, zeros(n-1))
    println("Error: ",et[1], " λ: ", λ)
    δq      = zeros(n-1)

    for iter = 1:maxiter # do a few iterations of the calibration

        ii = 1

        for p = 1:Ndatasets
            q = Q[p]
            N = size(q,1)
            Tfull = cat(3,[fkinePOE(xin,q[i,:]'+[δq;0]) for i = 1:N]...)
            Tm = squeeze(mean(Tfull,3),3)
            toOrthoNormal!(Tm)
            for i = 1:N
                # populate the matrices of the linear estimation problem
                y[ii:ii+5] = twistcoords(logT(Tm*trinv(Tfull[:,:,i])))
                A[ii:ii+5,:] = xii(q[i,1:(n-1)]'+δq, xin[:,1:(n-1)])
                ii += 6
            end
        end

        x = (A'A + λ*I)\A'y
        # Update the candidate nominal parameters
        δqi = δq + x

        et[iter+1] = evalErrorPOE_offsets_from_points(xin, Q, δqi)
        println("$iter Error: ",round(et[iter+1],5), " λ: ", λ, " Norm dx: ", round(norm(x),5))

        if norm(x) < 1e-10
            return δq, et
        end
        if et[iter+1] > et[iter]
            et[iter+1] = et[iter]
            λ *= 10
        else
            # λ = max(λ/10,1e-8)
            δq = deepcopy(δqi)
        end

    end

    return δq, et

end

function evalError(xin,Tn0,Ta,q)
    N = size(Ta,3)
    et = 0.0
    er = 0.0
    for i = 1:N
        Tfull = fkineLPOE(Tn0,xin,q[i,:]')
        et += norm(Ta[1:3,4,i]-Tfull[1:3,4])
        er += norm(skewcoords(logR(Ta[1:3,1:3,i]*Tfull[1:3,1:3]')))
    end
    return et/N, er/N
end

function evalErrorDual(xin,Tn0,q)
    N = size(Tn0,3)
    et = 0.0
    er = 0.0
    for i = 1:N
        T1 = fkineLPOE(Tn0[:,:,:,1],xin[:,:,1],q[i,:,1]')
        T2 = fkineLPOE(Tn0[:,:,:,2],xin[:,:,2],q[i,:,2]')
        et += norm(T2[1:3,4]-T1[1:3,4])
        er += norm(skewcoords(logR(T2[1:3,1:3]*T1[1:3,1:3]')))
    end
    return et/N, er/N
end

function evalErrorPOE(xin,Ta,q)
    xini = copy(xin)
    N = size(Ta,3)
    et = 0.0
    er = 0.0
    for i = 1:N
        Tfull = fkinePOE(xin,q[i,:]')
        et += norm(Ta[1:3,4,i]-Tfull[1:3,4])
        er += norm(skewcoords(logR(Ta[1:3,1:3,i]*Tfull[1:3,1:3]')))
    end
    return et/N, er/N
end

function evalErrorPOE_offsets_from_points(xi,Q,dq)
    Ndatasets = size(Q,1)
    J = Array{Float64}(Ndatasets)
    for p = 1:Ndatasets
        q = Q[p]
        N = size(q,1)
        T = Array{typeof(dq[1])}(4,4,N)
        for i = 1:N
            T[:,:,i] = fkinePOE(xi,q[i,:]' + [dq;0])
        end
        xyz = squeeze(T[1:3,4,:],2)
        mxyz = mean(xyz,2)
        Jpos = sum(abs((xyz.-mxyz)./mxyz))/N

        Rm = squeeze(mean(T[1:3,1:3,:],3),3)
        toOrthoNormal!(Rm)
        Jori = 0.0
        @inbounds for i = 1:N
            Jori += Rangle(T[1:3,1:3,i],Rm)
        end
        Jori /= Rangle(Rm)
        Jori /= N

        J[p] = Jpos + Jori
    end
    return sum(J)
end

function simulateCalibration1(N)
    n       = 6
    q       = 2π*rand(N,n)

    wn      = zeros(3,n)
    pn      = zeros(3,n)
    xin     = zeros(6,n+1)
    wn[:,1] = [0,0,1]
    pn[:,1] = [0,0,0]
    wn[:,2] = [0,1,0]
    pn[:,2] = [0,0,0.3]
    wn[:,3] = [0,1,0]
    pn[:,3] = [0,0,1]
    wn[:,4] = [1,0,0]
    pn[:,4] = [1,0,1]
    wn[:,5] = [0,1,0]
    pn[:,5] = [1,0,1]
    wn[:,6] = [1,0,0]
    pn[:,6] = [1,0,1]
    T0 = [  0 -1 0  1.2;
    0 0 -1 0;
    1 0  0  1;
    0 0  0  1]
    for i = 1:n
        xin[:,i] = [-skew(pn[:,i])*wn[:,i]; wn[:,i]]
    end
    xin[:,n+1] = twistcoords(logT(T0))
    xinmod   = deepcopy(xin)
    for i = 1:n+1
        xinmod[:,i]  += [0.001randn(3); 0.01π/180*randn(3)]
        #conformize(xinmod[:,i])
    end
    Ta  = zeros(4,4,N)
    for i = 1:N
        Ta[:,:,i] = fkinePOE(xin,q[i,:]')
    end
    return q, xin, T0, xinmod, Ta
end

function simulateCalibration2(N)
    n       = 6
    q       = 2π*rand(N,n)
    dh      = DH7600()
    AAA,BBB,T0,Ti0,Tn0 = jacobian(zeros(6),dh, eye(4));
    xin     = DH2twistsPOE(dh)
    xinmod   = deepcopy(xin)
    for i = 1:n+1
        xinmod[:,i]  += [0.001randn(3); 0.01π/180*randn(3)]
        #conformize(xinmod[:,i])
    end
    Ta  = zeros(4,4,N)
    for i = 1:N
        Ta[:,:,i] = fkinePOE(xin,q[i,:]')
    end
    return q, xin, T0, xinmod, Ta
end

function simulateCalibration3(N)
    srand(1)
    n       = 6
    q       = 2π*rand(N,n)
    dh      = DH7600()
    AAA,BBB,T0,Ti0,Tn0 = jacobian(zeros(6),dh, eye(4));
    xin = DH2twistsLPOE(Tn0)
    Tn0mod = deepcopy(Tn0)
    for i = 1:n+1
        Tn0mod[1:3,4,i] += 0.1randn(3)
    end
    Ta  = zeros(4,4,N)
    for i = 1:N
        Ta[:,:,i] = fkineLPOE(Tn0,xin,q[i,:]')
        AAA,BBB,T = jacobian(q[i,:]',dh, eye(4));
        @assert T ≈ Ta[:,:,i]
    end
    return q, xin, Tn0, Tn0mod, Ta
end

function simulateCalibration4(N)
    #srand(1)
    n       = 6
    qt      = 2π*rand(N,n)
    q1      = qt + 0.01*π/180*randn(N,n)
    q2      = qt + 0.01*π/180*randn(N,n)
    q       = cat(3,q1,q2)
    dh      = DH7600()
    AAA,BBB,T0,Ti0,Tn0 = jacobian(zeros(6),dh, eye(4));
    xin = DH2twistsLPOE(Tn0)
    xin = cat(3,xin,xin)
    Tn0mod = cat(4,deepcopy(Tn0),deepcopy(Tn0))
    for i = 1:n+1
        Tn0mod[1:3,4,i,1] += 0.1randn(3)
        Tn0mod[1:3,4,i,2] += 0.1randn(3)
    end
    Ta  = zeros(4,4,N)
    for i = 1:N
        Ta[:,:,i] = fkineLPOE(Tn0,xin,qt[i,:]')
    end

    return q, xin, Tn0, Tn0mod, Ta
end

# if false
#     pathopen    = "/work/fredrikb/flexifab/logs/calibrationNikon/2400_flexifabseamtracking.csv"
#     pathsave    = "/tmp/fredrikb/log.mat"
#     orcalog2mat(pathopen, pathsave)
#     ds          = 100
#     data        = MAT.matread(pathsave)
#     q           = getData("j_act_j*",data, ds, removeNaN = false)
#     q_frame     = getData("q_frame*",data, ds, removeNaN = false)
#
#     nans = any(isnan(q),2) | any(isnan(q_frame),2)
#
#     q = q[!nans,:]
#     q_frame = q_frame[!nans,:]
#     # plot(q)
#     using Quaternions
#
#     N = size(q,1)
#     Ta = zeros(4,4,N)
#     for i = 1:N
#         t = q_frame[i,6:8]'/1000
#         R = rotationmatrix(Quaternion(q_frame[i,1:4]...))
#         @assert isrot(R)
#         Ta[:,:,i] = [R t;0 0 0 1]
#     end
#
# end
if false
    q, xin,T0, xinmod, Ta= Robotlib.Calibration.simulateCalibration2(100)
    xic,et,er = calibPOE(xinmod,Ta,q,maxiter=150, λ = 10000.0)
    display(norm(xin[:,1:7]-xinmod[:,1:7]))
    display(norm(xin[:,1:7]-xic[:,1:7]))
end

if false
    q, xin,Tn0,Tn0mod, Ta = simulateCalibration3(100)
    @time Tn0c,xic,et,er = calibLPOE(xin,Tn0mod,Ta,q,maxiter=10, λ = 100.0)
    display(sqrt(sum((Tn0-Tn0mod).^2)/100))
    display(sqrt(sum((Tn0-Tn0c).^2)/100))
    display(round(Tn0-Tn0c,5))
end

if false
    N = 100
    q, xin,Tn0,Tn0mod, Ta = simulateCalibration4(N)
    @time Tn0c,xic,et,er = calibLPOEdual(xin,Tn0mod,q,maxiter=6, λ = 0.01)
    println("Error between Tn0 and Tn0mod: ",sqrt(sum((Tn0[1:3,4,:]-Tn0mod[1:3,4,:,1]).^2+(Tn0[1:3,4,:]-Tn0mod[1:3,4,:,2]).^2)/N))
    println("Error between Tn0 and Tn0c  : ",sqrt(sum((Tn0[1:3,4,:]-Tn0c[1:3,4,:,1]).^2+(Tn0[1:3,4,:]-Tn0c[1:3,4,:,2]).^2)/N))
    ei = 0.0
    ec = 0.0
    for i = 1:N
        T1 = fkineLPOE(Tn0mod,xin,q[i,:,1]')
        T2 = fkineLPOE(Tn0mod,xin,q[i,:,2]')
        ei += norm(twistcoords(logm(Ta[:,:,i]*trinv(T1))))
        ei += norm(twistcoords(logm(Ta[:,:,i]*trinv(T2))))
        T1 = fkineLPOE(Tn0c,xin,q[i,:,1]')
        T2 = fkineLPOE(Tn0c,xin,q[i,:,2]')
        ec += norm(twistcoords(logm(Ta[:,:,i]*trinv(T1))))
        ec += norm(twistcoords(logm(Ta[:,:,i]*trinv(T2))))
    end
    println("Initial error: ",round(ei/N,5), " Calibrated error: ", round(ec/N,5))
end

#plot(et[et .!= 0])


struct NDLCCRefSpace{T} <: RefSpace{T,6} end

function valuetype(ref::NDLCCRefSpace{T}, charttype::Type) where {T}
    SVector{universedimension(charttype),T}
end

function (ϕ::NDLCCRefSpace)(ndlc)

    u,v,w = parametric(ndlc)
    j = jacobian(ndlc)

    tu = tangents(ndlc,1)
    tv = tangents(ndlc,2)
    tw = tangents(ndlc,3)

    #B = [tv-tu tw-tu -tu]
    B = [tu tv tw]
    BMT = transpose(inv(B))

    du = SVector{3}(BMT[:,1])
    dv = SVector{3}(BMT[:,2])
    dw = SVector{3}(BMT[:,3])

    #Choose 1-(0,0,0), 2-(1,0,0), 3-(0,1,0), 4-(0,0,1)
    #Now it is listed as:
    #(1,2) [1-v-w,u,u]  (direction: from 1 to 2)
    #(1,3) [v,1-u-w,v]
    #(1,4) [w,w,1-u,v]
    #(2,3) [-v,u,0]
    #(2,4) [-w,0,u]
    #(3,4) [0,-w,v]

    # return SVector((
    #     (value=(BMT*[-v,u,0])       ,curl=(B*[0,0,2])/j) ,
    #     (value=(BMT*[-w,0,u])       ,curl=(B*[0,-2,0])/j),
    #     (value=(BMT*[(w+v-1),-u,-u]),curl=(B*[0,2,-2])/j),
    #     (value=(BMT*[0,-w,v])       ,curl=(B*[2,0,0])/j) ,
    #     (value=(BMT*[-v,(u+w-1),-v]),curl=(B*[-2,0,2])/j),
    #     (value=(BMT*[-w,-w,(u+v-1)]),curl=(B*[2,-2,0])/j)
    # ))

    return SVector((
        (value=-v*du + u*dv, curl=(2*tw)/j),
        (value=-w*du + u*dw, curl=(-2*tv)/j),
        (value=(w+v-1)*du - u*dv - u*dw, curl=(2*tv - 2*tw)/j),
        (value=-w*dv + v*dw, curl=(2*tu)/j),
        (value=-v*du + (u+w-1)*dv -v*dw, curl=(-2*tu + 2*tw)/j),
        (value=-w*du -w*dv + (u+v-1)*dw, curl=(2*tu -2*tv)/j),
    ))

##    	return SVector((
##        (value=(BMT*[(w+v-1),-u,-u]),curl=(B*[0,2,-2])/j),
##        (value=(BMT*[-v,(u+w-1),-v])  ,curl=(B*[-2,0,2])/j),
##        (value=(BMT*[-w,-w,(u+v-1)])  ,curl=(B*[2,-2,0])/j),
##        (value=(BMT*[-v,u,0])       ,curl=(B*[0,0,2])/j) ,
##        (value=(BMT*[-w,0,u])       ,curl=(B*[0,-2,0])/j),
##        (value=(BMT*[0,-w,v])       ,curl=(B*[2,0,0])/j)
##        ))
end

#check orientation
function curl(ref::NDLCCRefSpace, sh, el)
    a = [4,2,3,4,1,2]##[2,1,1,1,2,4]#[2,1,4,4,3,2]#[4,2,3,4,1,2]
    b = [3,4,2,1,3,1]##[3,2,4,3,4,3]#[4,3,3,1,2,1]#[3,4,2,1,3,1]
    sh1 = Shape(sh.cellid, b[sh.refid], -sh.coeff)
    sh2 = Shape(sh.cellid, a[sh.refid], sh.coeff)
    return [sh1,sh2]
end

function restrict(ϕ::NDLCCRefSpace{T}, dom1, dom2) where {T}
    # dom2 is the smaller of the domains

    K = numfunctions(ϕ)
    D = dimension(dom1)

    @assert K == 6
    @assert D == 3
    @assert D == dimension(dom2)

    Q = zeros(T,K,K)
    for (i,edge) in enumerate(CompScienceMeshes.edges(dom2))

        p = center(edge)
        c = cartesian(p)
        A = volume(edge)

        t = tangents(p,1)
        t = normalize(t)
        u = carttobary(dom1, c)
        x = neighborhood(dom1, u)

        y = ϕ(x)

        for j in 1:K
            Q[j,i] = -dot(y[j].value, t) * A
        end
    end

    return Q
end

function ttrace(x::NDLCCRefSpace, el, q, fc)
    # el: the supporting chart of the input
    # q: the local index of the face of el on which we compute the trace
    # fc: the chart on which we compute the trace

    el_ctr = cartesian(center(el))
    fc_ctr = cartesian(center(fc))

    T = scalartype(x)
    y = BEAST.RTRefSpace{T}()
    t = zeros(scalartype(x),3,6)
    for (j,edgej) in enumerate(CompScienceMeshes.edges(el))
        for (i,edgei) in enumerate(CompScienceMeshes.edges(fc))
            nbdi = center(edgei)
            nbdj = neighborhood(el, carttobary(el, cartesian(nbdi)))
            @assert isapprox(cartesian(nbdi), cartesian(nbdj), atol=1e-4)
            xval = x(nbdj)[j].value
            # yval = y(nbdi)[i].value
            tgt = -tangents(nbdi,1)
            nrm = normal(fc)
            bnr = cross(normalize(tgt), nrm)
            @assert dot(bnr, cartesian(nbdi)-fc_ctr) > 0
            @assert norm(bnr) ≈ 1
            dot(nrm, fc_ctr - el_ctr) < 0 && (nrm = -nrm)
            t[i,j] = dot(bnr, nrm × xval) * volume(edgei)
        end
    end
    return t
    #
    # # fa = SArray{Tuple{2},Int,1,2}
    # # te = SArray{Tuple{3},Int,1,3}
    # j = 0
    # y = BEAST.NDRefSpace{scalartype(x)}()
    # n = normal(fc)
    # for fac in faces(el)
    #     j += 1
    #     j == q && continue
    #     # fac != fc || continue
    #     #vind locale index (i) in fc van de edge die in zowel fc als fac ligt
    #     A = setdiff(fc.vertices,setdiff(fc.vertices,fac.vertices))
    #     # @show length(A)
    #     @assert length(A) == 2
    #     edg = fa([findfirst(isequal(A[1]), fc.vertices),findfirst(isequal(A[2]), fc.vertices)])
    #     i = abs(CompScienceMeshes.relorientation(edg,te([1,2,3])))
    #     #i = findfirst(isequal(setdiff(fcv,facv)[1]), fcv)
    #     edg_chart = simplex(fc.vertices[edg[1]], fc.vertices[edg[2]])
    #     edg_centr = cartesian(center(edg_chart))
    #     p = neighborhood(fc, carttobary(fc, edg_centr))
    #     r = neighborhood(el, carttobary(el, edg_centr))
    #     yp = y(p)[i].value
    #     xr = x(r)[j].value
    #     tgt = edg_chart.vertices[1] - edg_chart.vertices[2]
    #     ypt = dot(yp, tgt)
    #     xrt = dot(n × xr, tgt)
    #     # t[i,j] = volume(el)/(volume(fac)*norm(A[1]-A[2]))
    #     t[i,j] = xrt / ypt
    # end
    # return t
end