using SparseVectors
using Base.Test

import SparseVectors: exact_equal

### Data

rnd_x0 = sprand(50, 0.6)
rnd_x0f = full(rnd_x0)

rnd_x1 = sprand(50, 0.7) * 4.0
rnd_x1f = full(rnd_x1)

spv_x1 = SparseVector(8, [2, 5, 6], [1.25, -0.75, 3.5])
_x2 = SparseVector(8, [1, 2, 6, 7], [3.25, 4.0, -5.5, -6.0])
spv_x2 = view(_x2)

### Arithmetic operations

let x = spv_x1, x2 = spv_x2
    # negate
    @test exact_equal(-x, SparseVector(8, [2, 5, 6], [-1.25, 0.75, -3.5]))

    # abs and abs2
    @test exact_equal(abs(x), SparseVector(8, [2, 5, 6], abs([1.25, -0.75, 3.5])))
    @test exact_equal(abs2(x), SparseVector(8, [2, 5, 6], abs2([1.25, -0.75, 3.5])))

    # plus and minus
    xa = SparseVector(8, [1,2,5,6,7], [3.25,5.25,-0.75,-2.0,-6.0])

    @test exact_equal(x + x, x * 2)
    @test exact_equal(x + x2, xa)
    @test exact_equal(x2 + x, xa)

    xb = SparseVector(8, [1,2,5,6,7], [-3.25,-2.75,-0.75,9.0,6.0])
    @test exact_equal(x - x, SparseVector(8, Int[], Float64[]))
    @test exact_equal(x - x2, xb)
    @test exact_equal(x2 - x, -xb)

    @test full(x) + x2 == full(xa)
    @test full(x) - x2 == full(xb)
    @test x + full(x2) == full(xa)
    @test x - full(x2) == full(xb)

    # multiplies
    xm = SparseVector(8, [2, 6], [5.0, -19.25])
    @test exact_equal(x .* x, abs2(x))
    @test exact_equal(x .* x2, xm)
    @test exact_equal(x2 .* x, xm)

    @test full(x) .* x2 == full(xm)
    @test x .* full(x2) == full(xm)

    # max & min
    @test exact_equal(max(x, x), x)
    @test exact_equal(min(x, x), x)
    @test exact_equal(max(x, x2),
        SparseVector(8, Int[1, 2, 6], Float64[3.25, 4.0, 3.5]))
    @test exact_equal(min(x, x2),
        SparseVector(8, Int[2, 5, 6, 7], Float64[1.25, -0.75, -5.5, -6.0]))
end

### Complex

let x = spv_x1, x2 = spv_x2
    # complex
    @test exact_equal(complex(x, x),
        SparseVector(8, [2,5,6], [1.25+1.25im, -0.75-0.75im, 3.5+3.5im]))
    @test exact_equal(complex(x, x2),
        SparseVector(8, [1,2,5,6,7], [3.25im, 1.25+4.0im, -0.75+0.im, 3.5-5.5im, -6.0im]))
    @test exact_equal(complex(x2, x),
        SparseVector(8, [1,2,5,6,7], [3.25+0.im, 4.0+1.25im, -0.75im, -5.5+3.5im, -6.0+0.im]))

    # real & imag

    @test is(real(x), x)
    @test exact_equal(imag(x), sparsevector(Float64, length(x)))

    xcp = complex(x, x2)
    @test exact_equal(real(xcp), x)
    @test exact_equal(imag(xcp), x2)
end

### Zero-preserving math functions: sparse -> sparse

function check_nz2z_z2z{T}(f::Function, x::SparseVector{T}, xf::Vector{T})
    R = typeof(f(zero(T)))
    r = f(x)
    isa(r, AbstractSparseVector) || error("$f(x) is not a sparse vector.")
    eltype(r) == R || error("$f(x) results in eltype = $(eltype(r)), expect $R")
    all(r.nzval .!= 0) || error("$f(x) contains zeros in nzval.")
    full(r) == f(xf) || error("Incorrect results found in $f(x).")
end

for f in [floor, ceil, trunc, round]
    check_nz2z_z2z(f, rnd_x1, rnd_x1f)
end

for f in [log1p, expm1,
          sin, tan, sinpi, sind, tand,
          asin, atan, asind, atand,
          sinh, tanh, asinh, atanh]
    check_nz2z_z2z(f, rnd_x0, rnd_x0f)
end

### Non-zero-preserving math functions: sparse -> dense

function check_z2nz{T}(f::Function, x::SparseVector{T}, xf::Vector{T})
    R = typeof(f(zero(T)))
    r = f(x)
    isa(r, Vector) || error("$f(x) is not a dense vector.")
    eltype(r) == R || error("$f(x) results in eltype = $(eltype(r)), expect $R")
    r == f(xf) || error("Incorrect results found in $f(x).")
end

for f in [exp, exp2, exp10, log, log2, log10,
          cos, csc, cot, sec, cospi,
          cosd, cscd, cotd, secd,
          acos, acot, acosd, acotd,
          cosh, csch, coth, sech, acsch, asech]
    check_z2nz(f, rnd_x0, rnd_x0f)
end

### Reduction

# sum, sumabs, sumabs2, vecnorm

let x = spv_x1
    @test sum(x) == 4.0
    @test sumabs(x) == 5.5
    @test sumabs2(x) == 14.375

    @test vecnorm(x) == sqrt(14.375)
    @test vecnorm(x, 1) == 5.5
    @test vecnorm(x, 2) == sqrt(14.375)
    @test vecnorm(x, Inf) == 3.5
end

# maximum, minimum, maxabs, minabs

let x = spv_x1
    @test maximum(x) == 3.5
    @test minimum(x) == -0.75
    @test maxabs(x) == 3.5
    @test minabs(x) == 0.0
end

let x = abs(spv_x1)
    @test maximum(x) == 3.5
    @test minimum(x) == 0.0
end

let x = -abs(spv_x1)
    @test maximum(x) == 0.0
    @test minimum(x) == -3.5
end

let x = SparseVector(3, [1, 2, 3], [-4.5, 2.5, 3.5])
    @test maximum(x) == 3.5
    @test minimum(x) == -4.5
    @test maxabs(x) == 4.5
    @test minabs(x) == 2.5
end

let x = sparsevector(Float64, 8)
    @test maximum(x) == 0.0
    @test minimum(x) == 0.0
    @test maxabs(x) == 0.0
    @test minabs(x) == 0.0
end

# Issue https://github.com/JuliaLang/julia/issues/14046
let s14046 = sprand(5, 1.0)
    z = sparsevector(Float64, 5)
    @test z + s14046 == s14046
    @test 2*z == z + z == z
end


# Tests whether points project correctly

print("\tTest 1 - #generators == dim == 2...")
cone_1 = FiniteCone([[0.0, 1.0] [1.0, 1.0]])
test_data_1 = [([-1.0, 3.0], [0.0, 3.0]),
               ([-1.0, -1.0], [0.0, 0.0]),
               ([2.0, 0.0], [1.0, 1.0]),
               ([0.25, 1.0], [0.25, 1.0])]
for (x, proj_x) in test_data_1
    @test project(cone_1, x) ≈ proj_x
end
print("success!\n")

print("\tTest 2 - 3 = #generators > dim = 2...")
cone_2 = FiniteCone([[0.0, 1.0] [1.0, 1.0] [0.5, 1.0]])
test_data_2 = [([-1.0, 3.0], [0.0, 3.0]),
               ([-1.0, -1.0], [0.0, 0.0]),
               ([2.0, 0.0], [1.0, 1.0]),
               ([0.25, 1.0], [0.25, 1.0])]
for (x, proj_x) in test_data_2
    @test project(cone_2, x) ≈ proj_x
end
print("success!\n")

print("\tTest 3 - 2 = #generators < dim = 3...")
cone_3 = FiniteCone([[0.0, 0.0, 1.0] [1.0, 1.0, 1.0]])
test_data_3 = [([-1.0, -1.0, 1.0], [0.0, 0.0, 1.0]),
               ([3.0, 0.0, 0.0], [1.0, 1.0, 1.0]),
               ([-1.0, -1.0, -1.0], [0.0, 0.0, 0.0]),
               ([1.0, 2.0, 3.0], [1.5, 1.5, 3.0]),
               ([2.5, 2.5, 6.0], [2.5, 2.5, 6.0])]
for (x, proj_x) in test_data_3
    @test project(cone_3, x) ≈ proj_x
end
print("success!\n")


import HarmonicBalance: drop_powers, max_power
import HarmonicBalance.Symbolics.expand

@variables a,b,c

@test max_power(a^2 + b, a) == 2
@test max_power( a*((a+b)^4 )^2 + a, a) == 9

@test isequal(drop_powers(a^2+b, a, 1) , b)
@test isequal(drop_powers( (a+b)^2, a,1), b^2)
@test isequal(drop_powers( (a+b)^2, [a,b],1), 0)

@test isequal(drop_powers((a+b)^3 + (a+b)^5 , [a,b],4), expand((a+b)^3))


export iar

using IterativeSolvers
using LinearAlgebra
using Random
using Statistics

"""
    iar(nep,[maxit=30,][σ=0,][γ=1,][linsolvecreator=DefaultLinSolverCreator(),][tol=eps()*10000,][neigs=6,][errmeasure,][v=rand(size(nep,1),1),][logger=0,][check_error_every=1,][orthmethod=DGKS,][proj_solve=false,][inner_solver_method=DefaultInnerSolver(),][inner_logger=0])

Run the infinite Arnoldi method on the nonlinear eigenvalue problem stored in `nep`.

The target `σ` is the center around which eiganvalues are computed. The value `γ` corresponds
to scaling and specifying a shift and scaling is effectively the same as the transformation `λ=γs+σ` where `s`
is now the eigenvalue parameter. If you want eigenvalues in a disk centered, select `σ` as the center
of the disk and `γ` as the radius.
The vector `v` is the starting vector for constructing the Krylov space. The orthogonalization
method, used in contructing the orthogonal basis of the Krylov space, is specified by
`orthmethod`, see the package `IterativeSolvers.jl`.
The iteration is continued until `neigs` Ritz pairs have converged.
This function throws a `NoConvergenceException` if the wanted eigenpairs are not computed after `maxit` iterations.
However, if `neigs` is set to `Inf` the iteration is continued until `maxit` iterations without an error being thrown.
The parameter `proj_solve` determines if the Ritz paris are extracted using the Hessenberg matrix (false),
or as the solution to a projected problem (true). If `true`, the method is descided by `inner_solver_method`, and the
logging of the inner solvers are descided by `inner_logger`, which works in the same way as `logger`.

See [`augnewton`](@ref) for other parameters.

# Example
```julia-repl
julia> using NonlinearEigenproblems, LinearAlgebra
julia> nep=nep_gallery("dep0",100);
julia> v0=ones(size(nep,1));
julia> λ,v=iar(nep;v=v0,tol=1e-5,neigs=3);
julia> norm(compute_Mlincomb!(nep,λ[1],v[:,1])) # Is it an eigenvalue?
julia> λ    # print the computed eigenvalues
3-element Array{Complex{Float64},1}:
 -0.15606211475666945 - 0.12273439802763578im
 -0.15606211475666862 + 0.12273439802763489im
  0.23169243065648365 - 9.464790582509696e-17im
```

# References
* Algorithm 2 in Jarlebring, Michiels Meerbergen, A linear eigenvalue algorithm for the nonlinear eigenvalue problem, Numer. Math, 2012
"""
iar(nep::NEP;params...)=iar(ComplexF64,nep;params...)
function iar(
    ::Type{T},
    nep::NEP;
    orthmethod::Type{T_orth}=DGKS,
    maxit=30,
    linsolvercreator=DefaultLinSolverCreator(),
    tol=eps(real(T))*10000,
    neigs=6,
    errmeasure::ErrmeasureType = DefaultErrmeasure(nep),
    σ=zero(T),
    γ=one(T),
    v=randn(real(T),size(nep,1)),
    logger=0,
    check_error_every=1,
    proj_solve=false,
    inner_solver_method=DefaultInnerSolver(),
    inner_logger=0)where{T<:Number,T_orth<:IterativeSolvers.OrthogonalizationMethod}

    @parse_logger_param!(logger)
    @parse_logger_param!(inner_logger)

    # Ensure types σ and v are of type T
    σ=T(σ)
    v=Array{T,1}(v)

    n = size(nep,1);
    m = maxit;

    # initialization
    V = zeros(T,n*(m+1),m+1);
    H = zeros(T,m+1,m);
    y = zeros(T,n,m+1);
    α=Vector{T}(γ.^(0:m)); α[1]=zero(T);
    local M0inv::LinSolver=create_linsolver(linsolvercreator,nep,σ)

    err = NaN*ones(m,m);
    λ=zeros(T,m+1); Q=zeros(T,n,m+1);

    vv=view(V,1:1:n,1); # next vector V[:,k+1]
    vv[:]=v; vv[:]=vv[:]/norm(vv);
    k=1; conv_eig=0;
    local pnep::NEP;
    if (proj_solve)
        pnep=create_proj_NEP(nep);
    end

    while (k <= m) && (conv_eig<neigs)

        VV=view(V,1:1:n*(k+1),1:k); # extact subarrays, memory-CPU efficient
        vv=view(V,1:1:n*(k+1),k+1); # next vector V[:,k+1]

        y[:,2:k+1] = reshape(VV[1:1:n*k,k],n,k);
        broadcast!(/,view(y,:,2:k+1),view(y,:,2:k+1),(1:k)')
        y[:,1] = compute_Mlincomb!(nep,σ,y[:,1:k+1],α[1:k+1]);
        y[:,1] = -lin_solve(M0inv,y[:,1]);

        vv[:]=reshape(y[:,1:k+1],(k+1)*n,1);
        # orthogonalization
        H[k+1,k] = orthogonalize_and_normalize!(VV, vv, view(H,1:k,k), orthmethod)

        # compute Ritz pairs (every check_error_every iterations)
        if ((rem(k,check_error_every)==0)||(k==m))
            # Extract eigenvalues from Hessenberg matrix
            D,Z = eigen(H[1:k,1:k])

            VV = view(V,1:1:n,1:k)
            Q = VV*Z
            λ = σ .+ γ ./ D

            if (proj_solve)  # Projected solve to extract eigenvalues (otw hessenberg matrix)
                QQ,RR=qr(VV); # Project on this space
                QQ = Matrix(QQ)
                set_projectmatrices!(pnep,QQ,QQ);
                # Make a call to the inner solve method
                λproj,Qproj=inner_solve(inner_solver_method,T,pnep,
                                        V=RR*Z,λv=copy(λ),
                                        neigs=k,#size(λ,1)+3,
                                        σ=mean(λ),
                                        inner_logger=inner_logger);

                Q=QQ*Qproj;
                λ=λproj;
             end
            conv_eig=0;
            # compute the errors
            err[k,1:size(λ,1)]=
              map(s-> estimate_error(errmeasure,λ[s],Q[:,s]), 1:size(λ,1))
            # Log them and compute the converged
            push_iteration_info!(logger,2, k,err=err[k,1:size(λ,1)], λ=λ,
                                 continues=true);
            for s=1:size(λ,1)
                if err[k,s]<tol
                    conv_eig=conv_eig+1;
                    push_info!(logger,"+", continues=true);
                elseif err[k,s]<tol*10
                    push_info!(logger,"=", continues=true);
                else
                    push_info!(logger,"-", continues=true);
                end
            end
            push_info!(logger,"");
            # Sort the errors
            idx=sortperm(err[k,1:k]); # sort the error
            err[k,1:k]=err[k,idx];

            # extract the converged Ritzpairs
            if (k==m)||(conv_eig>=neigs)
                nrof_eigs = Int(min(length(λ),neigs))
                λ=λ[idx[1:nrof_eigs]]
                Q=Q[:,idx[1:length(λ)]]
            end

        end

        k=k+1;
    end
    k=k-1

    # NoConvergenceException
    if conv_eig<neigs && neigs != Inf
        # Sort the errors
        idx=sortperm(err[end,1:k]); err=err[end,idx];
        kk=length(λ);
        λ=λ[idx[1:kk]]; Q=Q[:,idx[1:kk]]; err=err[1:kk];
        msg="Number of iterations exceeded. maxit=$(maxit)."
        if conv_eig<3
            msg=string(msg, "Try to change the inner_solver_method for better performance.")
        end
        throw(NoConvergenceException(λ,Q,err,msg))
    end

    # extract the converged Ritzpairs
    λ=λ[1:min(length(λ),conv_eig)];
    Q=Q[:,1:min(size(Q,2),conv_eig)];
    return λ,Q,V[:,1:k]
end


using BenchmarkTools
using UpperEnvelope
using Plots
import Random

Random.seed!(42);

tmp1 = sortslices(reshape(rand(15), (5, 3)), dims=1)
k1=tmp1[:, 1]
v1=tmp1[:, 2]
w1=1 .+ tmp1[:, 3]

tmp2 = sortslices(reshape(rand(15), (5, 3)), dims=1)
k2=tmp2[:, 1]
v2=tmp2[:, 2]
w2=1 .+ tmp2[:, 3]

envelope=compute_envelope((k1, v1, w1), (k2, v2, w2), true)

# Plot example (too lazy to make code prettier)
p0=plot(k1, v1, marker=(:black), line=(:black), label="Function 1")
p1=plot(k1, v1, marker=(:black), line=(:black), label="Function 1")
plot!(p0, k2, v2, marker=(:x, :darkblue), line=(:dash, :darkblue), label="Function 2")
plot!(p1, k2, v2, marker=(:x, :darkblue), line=(:dash, :darkblue), label="Function 2")
plot!(p1,envelope[1], [envelope[2] envelope[3]], line=(:purple, 3), marker=(:diamond, :purple, 3), label="Upper envelope")
plot(p0, p1, size=0.8.*(980, 440), dpi=190)

savefig("example.png")


"""
    Polysegment <: AbstractImageBinarizationAlgorithm
    Polysegment()

    binarize([T,] img, f::Polysegment)
    binarize!([out,] img, f::Polysegment)

Uses the *polynomial segmentation* technique to group the image pixels
into two categories (foreground and background).

# Output

Return the binarized image as an `Array{Gray{T}}` of size `size(img)`. If
`T` is not specified, it is inferred from `out` and `img`.

# Details

The approach involves constructing a univariate second-degree polynomial
such that the two roots of the polynomial represent the graylevels
of two cluster centers (i.e the foreground and background). Pixels are then
assigned to the foreground or background depending on which cluster center
is closest.

# Arguments

The function argument is described in more detail below.

##  `img::AbstractArray`

The image that needs to be binarized. The image is automatically converted
to `Gray`.

# Example

Binarize the "cameraman" image in the `TestImages` package.

```julia
using TestImages, ImageBinarization

img = testimage("cameraman")
img_binary = binarize(img, Polysegment())
```

## Reference

1. R. E. Vidal, "Generalized Principal Component Analysis (GPCA): An Algebraic Geometric Approach to Subspace Clustering and Motion Segmentation." Order No. 3121739, University of California, Berkeley, Ann Arbor, 2003.
"""
struct Polysegment <: AbstractImageBinarizationAlgorithm end

function (f::Polysegment)(out::GenericGrayImage, img::GenericGrayImage)
  # Construct data matrix for second-degree polynomial (Equation 2.3) in [1].
  x = vec(img)
  Lₙ = hcat(ones(length(x)), x, x.^2)
  F = svd(Lₙ)
  c = F.Vt[end,:]
  p = Polynomial(vec(c))
  μ₁, μ₂ = sort(roots(p))
  # Binarize the image.
  @simd for i in CartesianIndices(img)
    val = img[i]
    out[i] = (val-μ₁)^2 < (val-μ₂)^2 ? 0 : 1
  end
  out
end

function (f::Polysegment)(out::GenericGrayImage, img::AbstractArray{<:Color3})
    # map `img` to grayspace while keeping the storage type
    # TODO: this mapping operation can be done lazily to reduce memory allocation
    GT = base_color_type(eltype(out)){eltype(eltype(img))}
    f(out, GT.(img))
end


julia> 0b1011010111 == 0o1327 == 0x2d7 == 727
true


"""
    iaaft(ts::AbstractArray{T, 1} where T;
            n_maxiter = 200, tol = 1e-6, n_windows = 50)

Generate an iteratively adjusted amplitude adjusted Fourier transform (IAAFT)
([Schreiber & Schmitz, 1996](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.77.635))
surrogate realization of `ts`.

## Arguments

- **`ts`**: the time series for which to generate an AAFT surrogate realization.

- **`n_maxiter`**: sets the maximum number of iterations to allow before ending
    the algorithm (if convergence is slow).

- **`tol`**: the relative tolerance for deciding if convergence is achieved.

- **`n_window`**: the number is windows used when binning the periodogram (used
    for determining convergence).

## References

T. Schreiber; A. Schmitz (1996). "Improved Surrogate Data for Nonlinearity
Tests". Phys. Rev. Lett. 77 (4): 635–638. doi:10.1103/PhysRevLett.77.635.
PMID 10062864. [https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.77.635](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.77.635)

"""
function iaaft(ts::AbstractVector{T} where T;
                n_maxiter = 200, tol = 1e-6, n_windows = 50)
    any(isnan.(ts)) && throw(DomainError(NaN,"The input must not contain NaN values."))

    # Sorted version of the original time series
    original_sorted = sort(ts)

    # Fourier transform and its amplitudes
    original_fft = fft(ts)
    original_fft_amplitudes = abs.(original_fft)

    # RANK ORDERING.
    # Create some Gaussian noise, and find the indices that sorts it with
    # increasing amplitude. Then sort the original time series according to
    # the indices rendering the Gaussian noise sorted.
    n = length(ts)
    g = rand(Normal(), n)
    inds = sortperm(g)
    ts_sorted = ts[inds]

    iter = 1
    success = false

    spectrum = Vector{Complex{Float64}}(undef, n)
    surrogate = Vector{Float64}(undef, n)

    diffs = zeros(Float64, 2)
    while iter <= n_maxiter
        # Take the Fourier transform of `ts_sorted` and get the phase angles of
        # the resulting complex numbers.
        FT = fft(ts_sorted)
        phase_angles = angle.(FT)

        # The new spectrum preserves the amplitudes of the Fourier transform of
        # the original time series, but randomises the phases (because the
        # phases are derived from the *randomly sorted* version of the original
        # time series).
        spectrum[:] = original_fft_amplitudes .* exp.(phase_angles .* 1im)

        # Now, let the surrogate time series be the values of the original time
        # series, but sorted according to the new spectrum. The shuffled series
        # is generated by taking the inverse Fourier transform of the spectrum
        # consisting of the original amplitudes, but having randomised phases.
        surrogate[:] = real(ifft(spectrum)) # ifft normalises by default

        # map original values onto shuffle
        surrogate[sortperm(surrogate)] = original_sorted

        ts_sorted[:] = surrogate

        # Convergence check
        periodogram = DSP.mt_pgram(ts)
        periodogram_surr = DSP.mt_pgram(surrogate)

        power_binned = interp([x for x in periodogram.freq],
                            periodogram.power,
                            n_windows)

        power_binned_surr = interp([x for x in periodogram_surr.freq],
                            periodogram_surr.power,
                            n_windows)

        if iter == 1
            diffs[1] = sum((power_binned[2] .- power_binned_surr[2]).^2) /
                        sum(power_binned[2].^2)
        else
            diffs[2] = sum((power_binned[2] .- power_binned_surr[2]).^2) /
                        sum(power_binned[2].^2)
            if abs(diffs[1] - diffs[2]) < tol
                break
            end
            diffs[1] = copy(diffs[2])
        end

        iter += 1
    end
    surrogate
end

"""
    iaaft_iters(ts::AbstractArray{T, 1} where T;
                n_maxiter = 100, tol = 1e-5, n_windows = 50)

Generate an iteratively adjusted amplitude adjusted Fourier transform (IAAFT) [1]
surrogate series for `ts` and return a vector containing the surrogate series
from each iteration. The last vector contains the final surrogate.

# Literature references
1. T. Schreiber; A. Schmitz (1996). "Improved Surrogate Data for Nonlinearity
Tests". Phys. Rev. Lett. 77 (4): 635–638. doi:10.1103/PhysRevLett.77.635. PMID
10062864.
"""
function iaaft_iters(ts::AbstractArray{T, 1} where T;
                        n_maxiter = 100, tol = 1e-5, n_windows = 50)

    # Sorted version of the original time series
    original_sorted = sort(ts)

    # Fourier transform and its amplitudes
    original_fft = fft(ts)
    original_fft_amplitudes = abs.(original_fft)

    # Create some Gaussian noise, and find the indices that sorts it with
    # increasing amplitude. Then sort the original time series according to the
    # indices rendering the Gaussian noise sorted.
    n = length(ts)
    g = rand(Normal(), n)
    inds = sortperm(g)
    ts_sorted = ts[inds]

    iter = 1
    success = false

    spectrum = Vector{Complex{Float64}}(undef, n)
    surrogate = Vector{Float64}(undef, n)
    surrogates = Vector{Vector{Float64}}(undef, 0)

    diffs = zeros(Float64, 2)
    while iter <= n_maxiter
        # Take the Fourier transform of `ts_sorted` and get the phase angles
        # of the resulting complex numbers.
        ts_sorted_fft = fft(ts_sorted)
        phase_angles = angle.(ts_sorted_fft)

        # The new spectrum preserves the amplitudes of the Fourier transform
        # of the original time series, but randomises the phases (because the
        # phases are derived from the *randomly sorted* version of the original
        # time series).
        spectrum .= original_fft_amplitudes .* exp.(phase_angles .* 1im)

        # Now, let the surrogate time series be the values of the original time
        # series, but sorted according to the new spectrum. The shuffled series
        # is generated by taking the inverse Fourier transform of the spectrum
        # consisting of the original amplitudes, but having randomised phases.
        surrogate[:] = real(ifft(spectrum)) # ifft normalises by default
        surrogate[sortperm(surrogate)] = original_sorted # map original values onto shuffle
        push!(surrogates, surrogate[:])

        ts_sorted .= surrogate

        # Convergence check
        periodogram = DSP.mt_pgram(ts)
        periodogram_surr = DSP.mt_pgram(surrogate)

        power_binned = interp([x for x in periodogram.freq],
                                periodogram.power,
                                n_windows)

        power_binned_surr = interp([x for x in periodogram_surr.freq],
                                    periodogram_surr.power,
                                    n_windows)

        if iter == 1
            diffs[1] = sum((power_binned[2] .- power_binned_surr[2]).^2) /
                        sum(power_binned[2].^2)
        else
            diffs[2] = sum((power_binned[2] .- power_binned_surr[2]).^2) /
                        sum(power_binned[2].^2)

            if abs(diffs[1] - diffs[2]) < tol
                break
            end
            diffs[1] = copy(diffs[2])
        end

        iter += 1
    end
    surrogates
end

export iaaft, iaaft_iters


function generictest_dixon_Fp(G, p = Characters.dixon_prime(G))
    F = Characters.FiniteFields.GF{p}
    ccG = Characters.conjugacy_classes(G)
    Ns = [Characters.CMMatrix(ccG, i) for i = 1:length(ccG)]
    @test isdiag(Characters.common_esd(Ns, F))
    esd = Characters.common_esd(Ns, F)

    let W = esd.basis
        @test inv(W) isa Matrix{F}
        @test isone(inv(W) * W)
        iW = inv(W)
        for N in Ns
            # basis actually diagonalizes all Ns
            m = W * N * iW
            @test isdiag(m)
        end
    end

    tbl = Characters.CharacterTable(F, G, ccG)
    @test Characters.irreducible_characters(tbl) isa
          Vector{<:Characters.Character}

    chars = Characters.irreducible_characters(tbl)

    # checking the degrees
    @test sum(Int.(SymbolicWedderburn.degree.(chars)) .^ 2) == order(G)

    # orthogonality of characters over F
    @test [dot(χ, ψ) for χ in chars, ψ in chars] == I
end

function generictest_dixon_C(G, p = Characters.dixon_prime(G))
    F = Characters.FiniteFields.GF{p}
    ccG = Characters.conjugacy_classes(G)
    tbl = Characters.CharacterTable(F, G, ccG)
    chars_Fp = Characters.irreducible_characters(tbl)

    degrees = SymbolicWedderburn.degree.(chars_Fp)

    m = Characters._multiplicities(chars_Fp)
    for i = 1:size(m, 1)
        @test all(Int.(m[i, :, :]) .<= degrees[i])
    end

    @test Characters.complex_character_table(Rational{Int}, tbl) isa
          Characters.CharacterTable{<:Group, <:Cyclotomic{Rational{Int}}}

    tblCC = Characters.complex_character_table(Rational{Int}, tbl)
    chars_CC = Characters.irreducible_characters(tblCC)

    id = one(G)
    @test [ψ(id) for ψ in chars_CC] == degrees

    # orthogonality of characters over ℂ
    @test [dot(χ, ψ) for χ in chars_CC, ψ in chars_CC] == I

    @test sum(χ -> SymbolicWedderburn.degree(χ)^2, chars_CC) == order(G)
end

@testset "Dixon Algorithm" begin
    @testset "DixonPrimes" begin
        @test Characters.dixon_prime(20, 20) == 41

        @testset "random" begin
            for n in rand(2:1000, 10)
                F = Characters.Primes.factor(n)
                e = lcm(collect(keys(F)))
                p = Characters.dixon_prime(n, e)
                @test mod(p, e) == 1
                @test p > 2 * floor(sqrt(n))
            end
        end

        @testset "DixonPrimeGroups" begin
            G = PermutationGroups.SymmetricGroup(4)
            ccG = Characters.conjugacy_classes(G)
            @test exponent(G) == 12
            @test exponent(ccG) == 12
            @test Characters.dixon_prime(G) ==
                  Characters.dixon_prime(ccG)
        end
    end

    @testset "Dixon over GF{p}" begin
        @testset "Example: Alt(4)" begin
            G = PermGroup([perm"(1,2,3)(4)", perm"(2,3,4)"])
            ccG = let S = gens(G)
                [
                    Orbit(S, one(G)),
                    Orbit(S, G(perm"(2,3,4)")),
                    Orbit(S, G(perm"(2,4,3)")),
                    Orbit(S, G(perm"(1,2)(3,4)")),
                ] # the order is taken from GAP
            end

            let ccG = ccG, p = 29
                F = Characters.FiniteFields.GF{p}
                Ns = [Characters.CMMatrix(ccG, i) for i = 1:length(ccG)]
                @test_throws AssertionError Characters.common_esd(Ns, F)
            end

            let ccG = ccG, p = 31
                F = Characters.FiniteFields.GF{p}
                Ns = [Characters.CMMatrix(ccG, i) for i = 1:length(ccG)]
                esd = Characters.EigenSpaceDecomposition(F.(Ns[1]))
                esd = Characters.refine(esd, F.(Ns[2]))
                esd = Characters.refine(esd, F.(Ns[3]))
                @test isdiag(esd)

                @test isdiag(Characters.common_esd(Ns, F))
            end

            # generic tests
            generictest_dixon_Fp(G)

            # tests specific to G
            p = Characters.dixon_prime(ccG)
            F = Characters.FiniteFields.GF{p}

            tbl = Characters.CharacterTable(F, G)
            chars = Characters.irreducible_characters(tbl)

            @test sort(SymbolicWedderburn.degree.(chars)) == [1, 1, 1, 3]

            @test Set(Int.(values(χ)) for χ in chars) ==
                  Set([[3, 0, 0, 6], [1, 4, 2, 1], [1, 2, 4, 1], [1, 1, 1, 1]])

        end
    end

    @testset "Lifting characters to ℂ" begin
        @testset "Example: Alt(4)" begin
            G = PermGroup([perm"(1,2,3)(4)", perm"(2,3,4)"])
            ccG = let S = gens(G)
                [
                    Orbit(S, one(G)),
                    Orbit(S, G(perm"(2,3,4)")),
                    Orbit(S, G(perm"(2,4,3)")),
                    Orbit(S, G(perm"(1,2)(3,4)")),
                ] # the order of cclasses is taken from GAP
            end

            @testset "PowerMap" begin
                pmG = Characters.PowerMap(ccG)
                @test size(pmG) == (4, 6)
                @test pmG[:, 0] == ones(Int, 4)
                @test pmG[:, 1] == 1:4
                @test pmG[:, 2] == [1, 3, 2, 1]
                @test pmG[:, 3] == [1, 1, 1, 4]
                @test pmG[:, 4] == [1, 2, 3, 1]
                @test pmG[:, 5] == [1, 3, 2, 4]
                @test all(isone.(first.(ccG) .^ 6))
            end

            # generic tests
            generictest_dixon_C(G)

            chars_C = Characters.irreducible_characters(G, ccG)
            E = Characters.Cyclotomics.E

            @test [collect(values(χ)) for χ in chars_C] == [
                E(3, 0) .* [3, 0, 0, -1],
                E(3, 0) .* [1, E(3, 2), E(3, 1), 1],
                E(3, 0) .* [1, E(3, 1), E(3, 2), 1],
                E(3, 0) .* [1, 1, 1, 1],
            ]
        end

        @testset "Example: Sym(4)" begin
            G = PermGroup([perm"(1,2)", perm"(1,2,3,4)"])
            ccG = let S = gens(G)
                ccG = [
                    Orbit(S, one(G)),
                    Orbit(S, G(perm"(1,2)(4)")),
                    Orbit(S, G(perm"(1,2)(3,4)")),
                    Orbit(S, G(perm"(1,2,3)(4)")),
                    Orbit(S, G(perm"(1,2,3,4)")),
                ] # the order of cclasses is taken from GAP
            end
            generictest_dixon_Fp(G)
            generictest_dixon_C(G)

            chars = Characters.irreducible_characters(G, ccG)

            @test sort(SymbolicWedderburn.degree.(chars)) == [1, 1, 2, 3, 3]
            @test [collect(values(χ)) for χ in chars] == [
                [2, 0, 2, -1, 0],
                [3, 1, -1, 0, -1],
                [1, 1, 1, 1, 1],
                [1, -1, 1, 1, -1],
                [3, -1, -1, 0, 1],
            ]
        end

        @testset "example: C₅ ⋊ C₄" begin
            G = PermGroup(perm"(1,2,4,5,3)", perm"(2,5,3,4)")
            S = gens(G)

            ccG = [
                Orbit(S, one(G)),
                Orbit(S, G(perm"(2,3)(4,5)")),
                Orbit(S, G(perm"(2,4,3,5)")),
                Orbit(S, G(perm"(2,5,3,4)")),
                Orbit(S, G(perm"(1,2,4,5,3)")),
            ]
            # the order of cclasses is taken from GAP
            generictest_dixon_Fp(G)
            generictest_dixon_C(G)

            chars = Characters.irreducible_characters(G, ccG)

            @test sort(SymbolicWedderburn.degree.(chars)) == [1, 1, 1, 1, 4]
            @test [collect(values(χ)) for χ in chars] == [
                E(4, 0) .* [4, 0, 0, 0, -1],
                E(4, 0) .* [1, 1, 1, 1, 1],
                E(4, 0) .* [1, 1, -1, -1, 1],
                E(4, 0) .* [1, -1, E(4), -E(4), 1],
                E(4, 0) .* [1, -1, -E(4), E(4), 1],
            ]
        end
    end

    @testset "Different base rings for characters_dixon" begin
        G = PermGroup(perm"(1,2,3,4)")
        @test eltype(Characters.irreducible_characters(Rational{Int}, G)) <:
            Characters.Character{<:Cyclotomic{Rational{Int}}}
        @test eltype(Characters.irreducible_characters(Rational{BigInt}, G)) <:
            Characters.Character{<:Cyclotomic{Rational{BigInt}}}
    end

    @time @testset "SmallPermGroups" begin
        for (ord, groups) in SmallPermGroups
            @testset "SmallGroup($ord, $n)" for (n, G) in enumerate(groups)
                @test Characters.irreducible_characters(G) isa
                      Vector{<:Characters.Character}
                generictest_dixon_Fp(G)
                generictest_dixon_C(G)
            end
        end
    end
end


using LinearAlgebra
using Random
using Test

n = 5
Xorig = randn(n, n)
A = UpperTriangular(rand(n, n))
B = randn(n, n)
B = B'*B
ws = LyapdWs(n)
for row = 1:n
    X = copy(Xorig)
    LinearRationalExpectations.solve_one_row!(X, A, B, n, row, ws)
    @test X[row, 1:row] ≈ -(A[row, row]*A[1:row, 1:row] - I(row))\B[row, 1:row]
    @test X[1:row, row] == X[row, 1:row]
end

n = 5
Xorig = randn(n, n)
A = zeros(n, n) .+ UpperTriangular(rand(n, n))
B = randn(n, n)
B = B'*B
ws = LyapdWs(n)
for row = 2:n
    A[row, row - 1] = 0.5
    X = copy(Xorig)
    Borig = copy(B)
    LinearRationalExpectations.solve_two_rows!(X, A, B, n, row, ws)
    @test B == Borig
    @test ws.AA2[1:2*row, 1:2*row] == kron(A[1:row, 1:row], A[(row - 1):row, (row - 1):row]) - I
    @test vec(X[(row-1):row, 1:row]) ≈ -(kron(A[1:row, 1:row], A[(row - 1):row, (row - 1):row]) - I(2*row))\vec(B[(row-1):row, 1:row])
    @test X[1:row, row-1] == X[row-1, 1:row]
    @test X[1:row, row] == X[row, 1:row]
    A[row, row - 1] = 0
end

for i = 1:100
    s = Int64(floor(1000*rand()))
    Random.seed!(s)
    local n = 5
    X1 = randn(n, n)
    X = X1'*X1
    local A = randn(n, n)
    F = eigen(A)
    if maximum(abs.(F.values)) > 1.0
        continue
    end
    local B = X - A*X*A'
    local ws = LyapdWs(n)
    Σ = similar(X)
    extended_lyapd!(Σ, A, B, ws)
    @test Σ ≈ X
end

# with unit roots
for i = 1:100
    s = Int64(floor(1000*rand()))
    Random.seed!(s)
    local n = 2
    X1 = randn(n, n)
    X = X1'*X1
    local A = randn(n, n)
    F = eigen(A)
    eigenvalues = diagm(rand(n))
    local A = real(F.vectors*eigenvalues*inv(F.vectors))
    m = 2
    AA = vcat(hcat([ 1.0 -0.5 0; 0 1.0 0], randn(m, n)),
              [1 -1 0 0 0],
              hcat(zeros(n, m+1), A))
    if abs(det(AA))  < 1e-12
        continue
    end
    local B = X - A*X*A'
    BB = vcat(hcat(I(m), zeros(m, n+1)),
              [0 0 0.2 0 0],
              hcat(zeros(n, m+1), B))

    local ws = LyapdWs(n + m + 1)
    Σ = Matrix{Float64}(undef, n+m+1, n+m+1)
    extended_lyapd!(Σ, AA, BB, ws)
    @test Σ[m + 1 .+ (1:n), m + 1 .+ (1:n)] ≈ X
end


"SPARK tableau for Gauss-Legendre Runge-Kutta method with s stages."
function TableauSPARKGLRK(s; name=Symbol("GLRK($s)"))
    g = TableauGauss(s)
    ω = get_GLRK_ω_matrix(s)
    δ = zeros(0, s)

    return TableauSPARK(name, g.o,
                        g.a, g.a, g.a, g.a,
                        g.a, g.a, g.a, g.a,
                        g.b, g.b, g.b, g.b,
                        g.c, g.c, g.c, g.b,
                        ω, δ)
end

"SPARK tableau for Gauss-Lobatto methods."
function TableauSPARKLobatto(A, B; name=Symbol("Lobatto($s)"))
    @assert A.s == B.s

    s = A.s
    o = min(A.o, B.o)
    δ = zeros(0, s)

    return TableauSPARK(name, o,
                        A.a, B.a, A.a, B.a,
                        A.a, B.a, A.a, B.a,
                        A.b, B.b, A.b, B.b,
                        A.c, B.c, A.c, A.b,
                        get_lobatto_ω_matrix(s), δ,
                        get_lobatto_nullvector(s))
end

"SPARK tableau for Gauss-Lobatto IIIA-IIIB method with s stages."
function TableauSPARKLobattoIIIAIIIB(s)
    TableauSPARKLobatto(TableauLobattoIIIA(s), TableauLobattoIIIB(s); name=Symbol("LobattoIIIAIIIB($s)"))
end

"SPARK tableau for Gauss-Lobatto IIIB-IIIA method with s stages."
function TableauSPARKLobattoIIIBIIIA(s)
    TableauSPARKLobatto(TableauLobattoIIIB(s), TableauLobattoIIIA(s); name=Symbol("LobattoIIIBIIIA($s)"))
end

"SPARK tableau for Gauss-Legendre/Gauss-Lobatto-IIIA-IIIB methods."
function TableauSPARKGLRKLobattoIIIAIIIB(s, σ=s+1)
    g = TableauGauss(s)
    A = TableauLobattoIIIA(σ)
    B = TableauLobattoIIIB(σ)

    α = A.a[2:σ, 1:σ]
    ã = B.a[1:σ, 1:s]

    o = min(g.o, A.o, B.o)
    δ = zeros(0, σ)

    return TableauSPARK(Symbol("GLRK($s)LobattoIIIAIIIB($σ)"), o,
                        g.a, g.a, α,   α,
                        ã,   ã,   A.a, B.a,
                        g.b, g.b, A.b, B.b,
                        g.c, g.c, A.c, A.b,
                        get_lobatto_ω_matrix(σ), δ)
end

"SPARK tableau for Gauss-Legendre/Gauss-Lobatto-IIIB-IIIA methods."
function TableauSPARKGLRKLobattoIIIBIIIA(s, σ=s+1)
    g = TableauGauss(s)
    A = TableauLobattoIIIA(σ)
    B = TableauLobattoIIIB(σ)

    α = B.a[2:σ, 1:σ]
    ã = A.a[1:σ, 1:s]

    o = min(g.o, B.o, A.o)
    δ = zeros(0, σ)

    return TableauSPARK(Symbol("GLRK($s)LobattoIIIAIIIB($σ)"), o,
                        g.a, g.a, α,   α,
                        ã,   ã,   B.a, A.a,
                        g.b, g.b, B.b, A.b,
                        g.c, g.c, A.c, A.b,
                        get_lobatto_ω_matrix(σ), δ)
end

function TableauSPARKLobatto(name, l1::Tableau{T}, l2::Tableau{T},
                                   l3::Tableau{T}, l4::Tableau{T}=l3, d=[]; R∞=1) where {T}

    @assert l1.s == l2.s == l3.s == l4.s

    ω = get_lobatto_ω_matrix(l3.s)
    δ = zeros(T, 0, 1)

    if length(d) == 0
        return TableauSPARK(name, min(l1.o, l2.o, l3.o, l4.o),
                            l3.a, l1.a, l4.a, l2.a,
                            l3.a, l1.a, l4.a, l2.a,
                            l3.b, l1.b, l4.b, l2.b,
                            l3.c, l1.c, l4.c, l4.b,
                            ω, δ)
    else
        @assert length(d) == q.s == p.s

        return TableauSPARK(name, min(l1.o, l2.o, l3.o, l4.o),
                            l3.a, l1.a, l4.a, l2.a,
                            l3.a, l1.a, l4.a, l2.a,
                            l3.b, l1.b, l4.b, l2.b,
                            l3.c, l1.c, l4.c, l4.b,
                            ω, δ, d)
    end
end

"Tableau for Gauss-Lobatto IIIA-IIIB-IIIC method with s stages."
function TableauSPARKLobABC(s)
    loba = TableauLobattoIIIA(s)
    lobb = TableauLobattoIIIB(s)
    lobc = TableauLobattoIIIC(s)
    TableauSPARKLobatto(Symbol("LobattoIIIABC($s)"), loba, lobc, lobb; R∞=(-1)^(s+1))
end

"Tableau for Gauss-Lobatto IIIA-IIIB-IIID method with s stages."
function TableauSPARKLobABD(s)
    loba = TableauLobattoIIIA(s)
    lobb = TableauLobattoIIIB(s)
    lobd = TableauLobattoIIID(s)
    TableauSPARKLobatto(Symbol("LobattoIIIABD($s)"), loba, lobd, lobb; R∞=(-1)^(s+1))
end

"SPARK Tableau for Variational Partitioned Runge-Kutta Methods."
function TableauSPARKVPRK(name, q::Tableau{T}, p::Tableau{T}, d=[]; R∞=1) where {T}

    @assert q.s == p.s

    ω = zeros(T, q.s, q.s+1)
    δ = zeros(T, 0, 1)

    for i in 1:q.s
        ω[i,i] = 1
    end

    if length(d) == 0
        return TableauSPARK(name, min(q.o, p.o),
                            q.a, p.a, q.a, p.a,
                            q.a, p.a, q.a, p.a,
                            q.b, p.b, q.b, p.b,
                            q.c, p.c, q.c, q.b,
                            ω, δ)
    else
        @assert length(d) == q.s == p.s

        return TableauSPARK(name, min(q.o, p.o),
                            q.a, p.a, q.a, p.a,
                            q.a, p.a, q.a, p.a,
                            q.b, p.b, q.b, p.b,
                            q.c, p.c, q.c, q.b,
                            ω, δ, d)
    end
end

"Tableau for Variational Gauss-Legendre method with s stages."
function TableauSPARKGLVPRK(s)
    glrk = TableauGauss(s)
    TableauSPARKVPRK(Symbol("GLVPRK($s)"), glrk, glrk; R∞=(-1)^s)
end


export kfactor

"""
`kfactor(G,k=1)` returns a `k`-factor of `G`. This is a set of edges
with the property that every vertex of `G` is incident with exactly `k`
edges in the set. An error is thrown if no such set exists.
"""
function kfactor(G::SimpleGraph, k::Int = 1)
    @assert k > 0 "The parameter k must be positive"

    VV = vlist(G)
    EE = elist(G)
    n = NV(G)
    m = NE(G)

    MOD = Model(get_solver())

    @variable(MOD, x[EE], Bin)

    for v in VV
        star = [get_edge(G, v, w) for w in G[v]]
        @constraint(MOD, sum(x[e] for e in star) == k)
    end

    optimize!(MOD)

    status = Int(termination_status(MOD))

    if status != 1
        error("This graph does not have a $k-factor")
    end

    X = value.(x)
    result = Set(e for e in EE if X[e] > 0.5)
    return result
end


### A Pluto.jl notebook ###
# v0.14.1

using Markdown
using InteractiveUtils

# ╔═╡ f643a980-5a13-11eb-1016-5f8ebc2b8370
using Plots

# ╔═╡ 13a311d0-5a11-11eb-1377-cd0b5a5ed7de
html"
The Lotka-Volterra Equations
"

# ╔═╡ 55724410-5a1f-11eb-2e9a-813c4e69c2b9
html"
Creating struct LotkaVolterra representing quantities in the equation:
"

# ╔═╡ 4f8f5730-5a11-11eb-1965-0d923a0baf77
@Base.kwdef mutable struct LotkaVolterra
	🐰::Float64	# Prey population
	🐺::Float64	# Predator population
	α::Float64 	# Prey growth factor per year
	β::Float64 	# Prey death factor per year
	δ::Float64	# Predator growth factor per year
	γ::Float64 	# Predator death factor per year
end

# ╔═╡ 839c59c0-5a1f-11eb-24c4-c7a3de982ca9
html"
Time step of 3x a month. The choice is arbitrary.
"

# ╔═╡ 13e4e050-5a12-11eb-1764-cf2089784c84
τ = 1/36

# ╔═╡ a1bc2480-5a1f-11eb-1238-cd6e072da32a
html"
The actual Lotka-Volterra differential equation:
"

# ╔═╡ 2bb5867e-5a12-11eb-135c-cf6728140506
function step!(l::LotkaVolterra)
	l.🐰 += τ*(l.α*l.🐰 - l.β*l.🐰*l.🐺)
	l.🐺 += τ*(l.δ*l.🐰*l.🐺 - l.γ*l.🐺)
end

# ╔═╡ b5717610-5a1f-11eb-03b8-5783009fa91f
html"
Setting things up for plotting:
"

# ╔═╡ 2ac5e370-5a1f-11eb-1966-e3f612893e56
begin
maxsteps = 6000
maxplots = 200
plotstep = maxsteps ÷ maxplots
xplot = zeros(maxplots)
yplot = zeros(maxplots)
tplot = zeros(maxplots)
end

# ╔═╡ c845dd80-5a1f-11eb-0ae1-41ce8fdd7be7
html"
Using real information about baboon/cheetah populations over time, we finally numerically solve the equations:
"

# ╔═╡ b144a830-5a12-11eb-3c23-9338902fef1e
let
lv = LotkaVolterra(🐰=10.0, 🐺=10.0, α=1.1, β=0.4, δ=0.4, γ=0.1)
iplot = 1
for step ∈ 1:maxsteps
	if step%plotstep == 1
		xplot[iplot] = lv.🐰
		yplot[iplot] = lv.🐺
		tplot[iplot] = τ*step
		iplot += 1
	end
	step!(lv)
end
end

# ╔═╡ fdc4e570-5a13-11eb-126a-df201e5b7370
let
plot(tplot, xplot, label="Baboon Population", lw=2)
plot!(tplot, yplot, label="Cheetah Population", lw=2)
plot!(title="Lotka-Volterra Equations", xlabel="Time (years)", ylabel="Population")
end

# ╔═╡ fd1f6030-5a1f-11eb-3c28-7d0a32456864
html"
Now, quickly, a less extreme example:
"

# ╔═╡ bcc7adad-03a4-49a9-9135-286fa3fc61f5
begin
steps = 720
plots = 180
pstep = steps ÷ plots
populations = 0.6:0.1:1.8
🐰 = zeros(length(populations), plots)
🐺 = zeros(length(populations), plots)
t = zeros(plots)
end

# ╔═╡ c19bd9c7-819b-437f-a628-30b2897d9958
let
for (i, pop) ∈ enumerate(populations)
	lk = LotkaVolterra(🐰=pop, 🐺=pop, α=2/3, β=4/3, δ=1, γ=1)
	iplot = 1
	for step ∈ 1:steps
		if step%pstep == 1
			🐰[i, iplot] = lk.🐰
			🐺[i, iplot] = lk.🐺
			t[iplot] = τ*step
			iplot += 1
		end
		step!(lk)
	end
end
end

# ╔═╡ 55b5d660-5a17-11eb-3db6-8394b4373f1b
let
plot(t, 🐰[1, :], label="Prey Population", lw=4)
plot!(t, 🐺[1, :], label="Predator Population", lw=4)
plot!(title="Lotka-Volterra Model",
		xlabel="Time (Years)", ylabel="Population (Thousands)")
end

# ╔═╡ 8950aca0-5a19-11eb-0e00-6f0f90f52a3f
let
plot(title = "Lotka-Volterra Phase Portrait")
xlabel!("Prey Population (Thousands)")
ylabel!("Predator Population (Thousands)")
for (i, pop) ∈ enumerate(populations)
	plot!(🐰[i, 1:100], 🐺[i, 1:100], label="x₀=y₀=$(pop)")
end
xₛ = 1/1 			# point of stability: xₛ = γ/δ = 1/1
yₛ = (2/3)/(4/3) 	# point of stability: yₛ = α/β = (2/3)/(4/3)
scatter!([xₛ], [yₛ], label="Equlibrium ($xₛ, $yₛ)")
end

# ╔═╡ Cell order:
# ╟─13a311d0-5a11-11eb-1377-cd0b5a5ed7de
# ╟─55724410-5a1f-11eb-2e9a-813c4e69c2b9
# ╠═4f8f5730-5a11-11eb-1965-0d923a0baf77
# ╟─839c59c0-5a1f-11eb-24c4-c7a3de982ca9
# ╠═13e4e050-5a12-11eb-1764-cf2089784c84
# ╟─a1bc2480-5a1f-11eb-1238-cd6e072da32a
# ╠═2bb5867e-5a12-11eb-135c-cf6728140506
# ╟─b5717610-5a1f-11eb-03b8-5783009fa91f
# ╠═2ac5e370-5a1f-11eb-1966-e3f612893e56
# ╟─c845dd80-5a1f-11eb-0ae1-41ce8fdd7be7
# ╠═b144a830-5a12-11eb-3c23-9338902fef1e
# ╠═f643a980-5a13-11eb-1016-5f8ebc2b8370
# ╠═fdc4e570-5a13-11eb-126a-df201e5b7370
# ╟─fd1f6030-5a1f-11eb-3c28-7d0a32456864
# ╠═bcc7adad-03a4-49a9-9135-286fa3fc61f5
# ╠═c19bd9c7-819b-437f-a628-30b2897d9958
# ╠═55b5d660-5a17-11eb-3db6-8394b4373f1b
# ╠═8950aca0-5a19-11eb-0e00-6f0f90f52a3f


using KNITRO, JuMP #, Base.Test

## Solve test problem 1 (Synthesis of processing system) in
 #  M. Duran & I.E. Grossmann, "An outer approximation algorithm for
 #  a class of mixed integer nonlinear programs", Mathematical
 #  Programming 36, pp. 307-339, 1986.  The problem also appears as
 #  problem synthes1 in the MacMINLP test set.
 #
 #  min   5 x4 + 6 x5 + 8 x6 + 10 x1 - 7 x3 -18 math.log(x2 + 1)
 #       - 19.2 math.log(x1 - x2 + 1) + 10
 #  s.t.  0.8 math.log(x2 + 1) + 0.96 math.log(x1 - x2 + 1) - 0.8 x3 >= 0
 #        math.log(x2 + 1) + 1.2 math.log(x1 - x2 + 1) - x3 - 2 x6 >= -2
 #        x2 - x1 <= 0
 #        x2 - 2 x4 <= 0
 #        x1 - x2 - 2 x5 <= 0
 #        x4 + x5 <= 1
 #        0 <= x1 <= 2
 #        0 <= x2 <= 2
 #        0 <= x3 <= 1
 #        x1, x2, x3 continuous
 #        x4, x5, x6 binary
 #
 #
 #  The solution is (1.30098, 0, 1, 0, 1, 0).
 ##

m = Model(solver=KnitroSolver(mip_method = KTR_MIP_METHOD_BB,
                              algorithm = KTR_ALG_ACT_CG,
                              outmode = KTR_OUTMODE_SCREEN,
                              KTR_PARAM_OUTLEV = KTR_OUTLEV_ALL,
                              KTR_PARAM_MIP_OUTINTERVAL = 1,
                              KTR_PARAM_MIP_MAXNODES = 10000,
                              KTR_PARAM_HESSIAN_NO_F = KTR_HESSIAN_NO_F_ALLOW))
x_U = [2,2,1]
@variable(m, x_U[i] >= x[i=1:3] >= 0)
@variable(m, y[4:6], Bin)

@NLobjective(m, Min, 10 + 10*x[1] - 7*x[3] + 5*y[4] + 6*y[5] + 8*y[6] - 18*log(x[2]+1) - 19.2*log(x[1]-x[2]+1))
@NLconstraints(m, begin
    0.8*log(x[2] + 1) + 0.96*log(x[1] - x[2] + 1) - 0.8*x[3] >= 0
    log(x[2] + 1) + 1.2*log(x[1] - x[2] + 1) - x[3] - 2*y[6] >= -2
    x[2] - x[1] <= 0
    x[2] - 2*y[4] <= 0
    x[1] - x[2] - 2*y[5] <= 0
    y[4] + y[5] <= 1
end)
solve(m)

ktrmod = internalmodel(m)
MathProgBase.freemodel!(ktrmod)


using Plots
using CSV
using Printf
using DataFrames
using LaTeXStrings
using LsqFit            # Required for the statistics on the fitting process
using FrictionModels

# Read in a CSV file containing the data to fit.
#
# The file is structured as follows:
# Column 1: Time
# Column 2: Velocity
# Column 3: Normal force
# Column 4: Friction Force
println(pwd())
data = CSV.read("examples//test_data_1.csv", DataFrame)

# Construct an initial estimate for model parameters
init_guess = CoulombModel(
    0.25
)

# Fit the model
results = fit_model(
    init_guess,     # Initial Guess
    data[:,1],      # Time
    data[:,4],      # Friction Force
    data[:,3],      # Normal Force
    data[:,2]       # Velocity
)

# Solve the model
F = friction.(results.model, data[:,3], data[:,2])

# Plot the data
plt = plot(
    data[:,2], F,
    xlabel = L"v(t)",
    ylabel = L"F(t)",
    label = "Fitted Model",
    lw = 2,
    legend = :bottomright
)
plot!(
    plt,
    data[:,2], data[:,4],
    st = :scatter,
    label = "Measured Data"
)
display(plt)

# Print out the coefficients of the fitted model
@printf("Fitted Model Coefficients:\n")

# Print out statistitics related to the fit
sigma = stderror(results.fit)
confidence = confidence_interval(results.fit)

# Print out the results
@printf(
    "coefficient: %f\n\tError: %f\n\tConfidence Interval:: (%f, %f)\n",
    results.model.coefficient,
    sigma[1],
    confidence[1][1],
    confidence[1][2]
)


function model_uptake(
    conv::Function = expconv;
    t::AbstractVector,
    params::NamedTuple,
    ca::AbstractVector,
)
    @extract (fp, ps, vp) params
    Tp = vp / (fp + ps)
    E = ps / (fp + ps)
    ca_conv_1 = zeros(length(t))
    for idx = 2:length(ca)
        ca_conv_1[idx] = ca_conv_1[idx-1] + (t[idx] - t[idx-1]) * ca[idx]
    end
    ct = fp .* ((1 - E) .* expconv(ca, 1 / Tp, t) .+ E .* ca_conv_1)
    return ct
end

function fit_uptake_nls(
    conv::Function = expconv;
    t::AbstractVector,
    ca::AbstractVector,
    ct::AbstractArray,
    mask = true,
)
    (t, ct, mask, num_timepoints, volume_size) = resolve_fitting_inputs(; t, ca, ct, mask)
    fp, ps, vp = (fill(NaN, volume_size...) for _ = 1:3)
    model(x, p) = model_uptake(; t = x, ca, params = (fp = p[1], ps = p[2], vp = p[3]))
    lls_est = fit_uptake_lls(; t, ca, ct, mask).est
    init_fp, init_ps, init_vp = select(lls_est, (:fp, :ps, :vp))
    for idx in eachindex(IndexCartesian(), mask)
        if mask[idx] == false
            continue
        end
        initialvalue = [init_fp[idx], init_ps[idx], init_vp[idx]]
        (fp[idx], ps[idx], vp[idx]) = curve_fit(model, t, ct[idx, :], initialvalue).param
    end
    return (; est = (; fp, ps, vp))
end

function fit_uptake_lls(;
    t::AbstractVector,
    ca::AbstractVector,
    ct::AbstractArray,
    mask = true,
)
    (t, ct, mask, num_timepoints, volume_size) = resolve_fitting_inputs(; t, ca, ct, mask)
    x1, fp, x3, ps, vp = (fill(NaN, volume_size...) for _ = 1:5)

    M = zeros(num_timepoints, 3)
    M[:, 2] = cumul_integrate(t, ca, TrapezoidalFast())
    M[:, 3] = cumul_integrate(t, M[:, 2], TrapezoidalFast())
    for idx in eachindex(IndexCartesian(), mask)
        if mask[idx] == false
            continue
        end
        M[:, 1] = -cumul_integrate(t, ct[idx, :], TrapezoidalFast())
        (x1[idx], fp[idx], x3[idx]) = M \ ct[idx, :]
    end
    denum = @. x1 * fp - x3
    @. ps = fp * x3 / denum
    @. vp = fp^2 / denum
    @. ps[denum==0] = 0
    @. vp[denum==0] = 0
    return (; est = (; fp, ps, vp))
end


function part1(input)
    binaries = map(line -> collect(line) .== '1', readlines(input))
    majority = sum(binaries) .> length(binaries) / 2
    γ = decode(majority)
    ϵ = decode(.!majority)
    return γ * ϵ
end

function part2(input)
    binaries = map(line -> collect(line) .== '1', readlines(input))
    oxygen = decode(find_common(copy(binaries), 1, false))
    CO2 = decode(find_common(copy(binaries), 1, true))
    return oxygen * CO2
end

decode(x) = sum(x .<< (reverse(eachindex(x)) .- 1))

function find_common(binaries, n, choose_uncommon)
    length(binaries) == 1 && return only(binaries)
    common = sum(getindex.(binaries, n)) >= length(binaries) / 2
    return find_common(filter(x -> x[n] == common ⊻ choose_uncommon, binaries),
                       n + 1, choose_uncommon)
end


abstract type AbstractEquationOfState <: AbstractPhysicsComponent end

Base.@kwdef struct BarotropicFluid{FT} <: AbstractEquationOfState
    ρₒ::FT
    cₛ::FT
end

Base.@kwdef struct DryIdealGas{FT} <: AbstractEquationOfState
    R::FT
    pₒ::FT
    γ::FT
end

# Energy Prognostic Cases, perhaps abstract type

Base.@kwdef struct TotalEnergy{FT} <: AbstractEquationOfState
    γ::FT
end

Base.@kwdef struct DryEuler{FT} <: AbstractEquationOfState
    γ::FT
end

Base.@kwdef struct LinearizedTotalEnergy{FT} <: AbstractEquationOfState
    γ::FT
end

Base.@kwdef struct LinearizedDryEuler{FT} <: AbstractEquationOfState
    γ::FT
end

# Linearized Equation of States
linearize(eos::TotalEnergy) = LinearizedTotalEnergy(eos.γ)
linearize(eos::DryEuler) = LinearizedDryEuler(eos.γ)

"""
  Thermodynamic relationships
"""

@inline function calc_pressure(eos::BarotropicFluid, state)
    cₛ = eos.cₛ
    ρₒ = eos.ρₒ
    ρ = state.ρ

    return (cₛ * ρ)^2 / (2 * ρₒ)
end

@inline function calc_sound_speed(eos::BarotropicFluid, state)
    cₛ = eos.cₛ
    ρₒ = eos.ρₒ
    ρ = state.ρ

    return cₛ * sqrt(ρ / ρₒ)
end

@inline function calc_pressure(eos::DryIdealGas, state)
    R = eos.R
    pₒ = eos.pₒ
    γ = eos.γ
    ρθ = state.ρθ

    return pₒ * (R / pₒ * ρθ)^γ
end

@inline function calc_sound_speed(eos::DryIdealGas, state)
    R = eos.R
    pₒ = eos.pₒ
    γ = eos.γ
    ρ = state.ρ

    return sqrt(γ * calc_pressure(eos, state) / ρ)
end

"""
Modified Maciek's world
"""
@inline function calc_pressure(eos::TotalEnergy, state, aux)
    γ = eos.γ
    Φ = aux.Φ
    ρ = state.ρ
    ρe = state.ρe
    ρu = state.ρu
    return (γ - 1) * (ρe - dot(ρu, ρu) / 2ρ - ρ * Φ)
end

@inline function calc_sound_speed(eos::TotalEnergy, state, aux)
    γ = eos.γ
    ρ = state.ρ
    return sqrt(γ * calc_pressure(eos, state, aux) / ρ)
end

@inline function calc_pressure(eos::LinearizedTotalEnergy, state, aux)
    γ = eos.γ
    Φ = aux.Φ
    ρ = state.ρ
    ρe = state.ρe
    return (γ - 1) * (ρe - ρ * Φ)
end

@inline function calc_sound_speed(eos::LinearizedTotalEnergy, state, aux)
    γ = eos.γ
    ρ = state.ρ
    return sqrt(γ * calc_pressure(eos, state, aux) / ρ)
end

@inline function calc_pressure(eos::DryEuler, state, aux)
    γ = eos.γ
    ρ = state.ρ
    ρe = state.ρe
    ρu = state.ρu
    return (γ - 1) * (ρe - dot(ρu, ρu) / 2ρ)
end

@inline function calc_sound_speed(eos::DryEuler, state, aux)
    γ = eos.γ
    ρ = state.ρ
    return sqrt(γ * calc_pressure(eos, state, aux) / ρ)
end

@inline function calc_pressure(eos::LinearizedDryEuler, state, aux)
    γ = eos.γ
    ρe = state.ρe
    return (γ - 1) * ρe
end

@inline function calc_sound_speed(eos::LinearizedDryEuler, state, aux)
    γ = eos.γ
    ρ = state.ρ
    return sqrt(γ * calc_pressure(eos, state, aux) / ρ)
end

"""
  Maciek's world
"""
function pressure(ρ, ρu, ρe, Φ)
    FT = eltype(ρ)
    γ = FT(gamma(param_set))
    if total_energy
        (γ - 1) * (ρe - dot(ρu, ρu) / 2ρ - ρ * Φ)
    else
        (γ - 1) * (ρe - dot(ρu, ρu) / 2ρ)
    end
end

function totalenergy(ρ, ρu, p, Φ)
    FT = eltype(ρ)
    γ = FT(gamma(param_set))
    if total_energy
        return p / (γ - 1) + dot(ρu, ρu) / 2ρ + ρ * Φ
    else
        return p / (γ - 1) + dot(ρu, ρu) / 2ρ
    end
end

function soundspeed(ρ, p)
    FT = eltype(ρ)
    γ = FT(gamma(param_set))
    sqrt(γ * p / ρ)
end

@inline function linearized_pressure(ρ, ρe, Φ)
    FT = eltype(ρ)
    γ = FT(gamma(param_set))
    if total_energy
        (γ - 1) * (ρe - ρ * Φ)
    else
        (γ - 1) * ρe
    end
end

"""
  Base extensions
"""
function info(::AbstractEquationOfState)
    error("Not implemented!")
end

function info(::BarotropicFluid)
    println("The equation of state (eos) is:")
    printstyled("eos = (cₛ ρ)^2 / (2 ρₒ) \n", color = 82)
    println("The sound speed is:")
    printstyled("soundspeed = cₛ sqrt(ρ / ρₒ)  \n", color = 82)
    println("ρ : density ")
    println("cₛ: reference soundspeed ")
    println("ρₒ: reference density ")
end

function info(::DryIdealGas)
    error("Not implemented!")
end


export DFT, IDFT

abstract type FourierTransform{N,C,D,T1,T2} <: LinearOperator end

"""
`DFT([domainType=Float64::Type,] dim_in::Tuple [,dims])`

`DFT(dim_in...)`

`DFT(x::AbstractArray [,dims])`

Creates a `LinearOperator` which, when multiplied with an array `x::AbstractArray{N}`, returns the `N`-dimensional Discrete Fourier Transform over dimensions `dims` of `x`.

```julia
julia> DFT(Complex{Float64},(10,10))
ℱ  ℂ^(10, 10) -> ℂ^(10, 10)

julia> DFT(10,10)
ℱ  ℝ^(10, 10) -> ℂ^(10, 10)

julia> A = DFT(ones(3))
ℱ  ℝ^3 -> ℂ^3

julia> A*ones(3)
3-element Array{Complex{Float64},1}:
 3.0+0.0im
 0.0+0.0im
 0.0+0.0im
```

"""
struct DFT{N,
	      C<:RealOrComplex,
	      D<:RealOrComplex,
	      T1<:AbstractFFTs.Plan,
	      T2<:AbstractFFTs.Plan} <: FourierTransform{N,C,D,T1,T2}
	dim_in::NTuple{N,Int}
	A::T1
	At::T2
end

"""
`IDFT([domainType=Float64::Type,] dim_in::Tuple [,dims])`

`IDFT(dim_in...)`

`IDFT(x::AbstractArray [,dims])`

Creates a `LinearOperator` which, when multiplied with an array `x::AbstractArray{N}`, returns the `N`-dimensional Inverse Discrete Fourier Transform over dimensions `dims` of `x`.

```julia
julia> IDFT(Complex{Float64},(10,10))
ℱ⁻¹  ℂ^(10, 10) -> ℂ^(10, 10)

julia> IDFT(10,10)
ℱ⁻¹ ℝ^(10, 10) -> ℂ^(10, 10)

julia> A = IDFT(ones(3))
ℱ⁻¹  ℝ^3 -> ℂ^3

julia> A*ones(3)
3-element Array{Complex{Float64},1}:
 1.0+0.0im
 0.0+0.0im
 0.0+0.0im
```

"""
struct IDFT{N,
	       C<:RealOrComplex,
	       D<:RealOrComplex,
	       T1<:AbstractFFTs.Plan,
	       T2<:AbstractFFTs.Plan} <: FourierTransform{N,C,D,T1,T2}
	dim_in::NTuple{N,Int}
	A::T1
	At::T2
end

# Constructors
#standard constructor
DFT(dim_in::NTuple{N,Int},dims=1:N) where {N} = DFT(zeros(dim_in),dims)

function DFT(x::AbstractArray{D,N},dims=1:ndims(x)) where {N,D<:Real}
	A,At = plan_fft(x,dims), plan_bfft(fft(x,dims),dims)
	DFT{N,Complex{D},D,typeof(A),typeof(At)}(size(x),A,At)
end

function DFT(x::AbstractArray{D,N},dims=1:ndims(x)) where {N,D<:Complex}
	A,At = plan_fft(x,dims), plan_bfft(fft(x,dims),dims)
	DFT{N,D,D,typeof(A),typeof(At)}(size(x),A,At)
end

DFT(T::Type,dim_in::NTuple{N,Int},dims=1:N) where {N} = DFT(zeros(T,dim_in),dims)
DFT(dim_in::Vararg{Int}) = DFT(dim_in)
DFT(T::Type,dim_in::Vararg{Int}) = DFT(T,dim_in)

function IDFT(x::AbstractArray{D,N},dims=1:ndims(x)) where {N,D<:Real}
	A,At = plan_ifft(x,dims), plan_fft(ifft(x,dims),dims)
	IDFT{N,Complex{D},D,typeof(A),typeof(At)}(size(x),A,At)
end

#standard constructor
IDFT(T::Type,dim_in::NTuple{N,Int},dims=1:N) where {N} = IDFT(zeros(T,dim_in),dims)

function IDFT(x::AbstractArray{D,N},dims=1:ndims(x)) where {N,D<:Complex}
	A,At = plan_ifft(x,dims), plan_fft(ifft(x,dims),dims)
	IDFT{N,D,D,typeof(A),typeof(At)}(size(x),A,At)
end

IDFT(dim_in::NTuple{N,Int},dims=1:N) where {N} = IDFT(zeros(dim_in),dims)
IDFT(dim_in::Vararg{Int}) = IDFT(dim_in)
IDFT(T::Type,dim_in::Vararg{Int}) = IDFT(T,dim_in)

# Mappings

function mul!(y::AbstractArray{C,N},
              L::DFT{N,C,C,T1,T2},
              b::AbstractArray{C,N}) where {N,C<:Complex,T1,T2}
	mul!(y,L.A,b)
end

function mul!(y::AbstractArray{C,N},
              L::DFT{N,C,D,T1,T2},
              b::AbstractArray{D,N}) where {N,C<:Complex,D<:Real,T1,T2}
	mul!(y,L.A,complex(b))
end

function mul!(y::AbstractArray{C,N},
              L::AdjointOperator{DFT{N,C,C,T1,T2}},
              b::AbstractArray{C,N}) where {N,C<:Complex,T1,T2}
	mul!(y,L.A.At,b)
end

function mul!(y::AbstractArray{D,N},
              L::AdjointOperator{DFT{N,C,D,T1,T2}},
              b::AbstractArray{C,N}) where {N,C<:Complex,D<:Real,T1,T2}
	y2 = complex(y)
	mul!(y2,L.A.At,b)
	y .= real.(y2)
end

function mul!(y::AbstractArray{C,N},
              L::IDFT{N,C,C,T1,T2},
              b::AbstractArray{C,N}) where {N,C<:Complex,T1,T2}
	mul!(y,L.A,b)
end

function mul!(y::AbstractArray{C,N},
              L::IDFT{N,C,D,T1,T2},
              b::AbstractArray{D,N}) where {N,C<:Complex,D<:Real,T1,T2}
	mul!(y,L.A,complex(b))
end

function mul!(y::AbstractArray{C,N},
              L::AdjointOperator{IDFT{N,C,C,T1,T2}},
              b::AbstractArray{C,N}) where {N,C<:Complex,T1,T2}
	mul!(y,L.A.At,b)
	y .*= L.A.A.scale
end

function mul!(y::AbstractArray{D,N},
              L::AdjointOperator{IDFT{N,C,D,T1,T2}},
              b::AbstractArray{C,N}) where {N,C<:Complex,D<:Real,T1,T2}

	y2 = complex(y)
	mul!(y2,L.A.At,b)
	y .= real.(y2).*L.A.A.scale
end

# Properties

size(L::FourierTransform) = (L.dim_in,L.dim_in)

fun_name(A::DFT) = "ℱ"
fun_name(A::IDFT) = "ℱ⁻¹"

domainType(L::FourierTransform{N,C,D,T1,T2}) where {N,C,D,T1,T2} = D
codomainType(L::FourierTransform{N,C,D,T1,T2}) where {N,C,D,T1,T2} = C

is_AcA_diagonal(L::FourierTransform)     = true
is_AAc_diagonal(L::FourierTransform)     = true
is_invertible(L::FourierTransform)       = true
is_full_row_rank(L::FourierTransform)    = true
is_full_column_rank(L::FourierTransform) = true

diag_AcA(L::DFT) = float(prod(size(L,1)))
diag_AAc(L::DFT) = float(prod(size(L,1)))

diag_AcA(L::IDFT) = 1/prod(size(L,1))
diag_AAc(L::IDFT) = 1/prod(size(L,1))


using Random,StatsBase,Distributions,Polynomials
Random.seed!(150000)

#LSM原始算法，针对美式PUT
function LSM_PRIMAL(S0::Float64,K::Float64,T::Float64,r::Float64,σ::Float64,I::Int,M::Int)
    """

    # Arguments
    - `S0`: 初始标的值.
    - `K`: 行权值
    - `T`: 到期时间，年为单位
    - `r`: 无风险利率
    - `σ`: 波动率
    - `I`: 原始算法中路径数
    - `M`: 时间步数

    """

    dt=T/M
    df=exp(-r*dt)

    p1=(r-0.5*(σ^2))*dt
    p2=rand(Normal(0,1),M+1,I)
    p3=p1 .+σ*√(dt)*p2
    p4=cumsum(p3,dims=1)

    S=S0*(p4.|>exp)
    S[1,:]=repeat([S0],I)
    h=map(x->max(K-x,0),S)

    V=copy(h[end,:])

    @simd for t=M:-1:1
        rg=polyfit(S[t,:],V*df,5)
        C=map(x->polyval(rg,x),S[t,:])
        res=h[t,:].>C
        for i=1:I
            if res[i]==true
                V[i]=h[t,i]
            else
                V[i]=V[i]*df
            end
        end
    end

    v0=df*sum(V)/I

end

#LSM对偶算法
function LSM_DUAL(S0::Float64,K::Float64,T::Float64,r::Float64,σ::Float64,J::Int,I1::Int,I2::Int,M::Int,reg::Int)
    """

    # Arguments
    - `S0`: 初始::Int标的值::Int.
    ::Int- ::Int`K`: 行权值
    - `T`: 到期时间，年为单位
    - `r`: 无风险利率
    - `σ`: 波动率
    - `J`: 对偶算法中子循环模拟路径
    - `I1`: 原始算法中路径数
    - `I2`: 对偶算法中路径数
    - `M`: 时间步数
    - `reg`: 多项是回归阶数

    """
    dt=T/M
    df=exp(-r*dt)

    #路径I1,产生回归系数
    S1=genrate_path(S0,r,σ,dt,M,I1)
    h1=map(x->max(K-x,0),S1)
    V1=map(x->max(K-x,0),S1)

    #路径I1,产生回归系数
    rg0=[]
    a=Poly(0)
    push!(rg0,a)
    @simd for t=M:-1:2
        rgr=polyfit(S1[t,:],V1[t+1,:]*df,5)
        push!(rg0,rgr)
    end
    push!(rg0,a)
    rg=reverse(rg0)

    #路径I2，对偶算法
    Q=zeros(M+1,I2)
    U=zeros(M+1,I2)
    S=genrate_path(S0,r,σ,dt,M,I2)
    h=map(x->max(K-x,0),S)
    V=map(x->max(K-x,0),S)

    @simd for t=2:M+1
        for i=1:I2
            tmp=polyval(rg[t],S[t,i])
            Vt=max(h[t,i],tmp)
            St=generate_nest_mc(r,σ,dt,S[t-1,i],J)
            Ct=map(x->polyval(rg[t],x),St)
            ht=map(x->max(K-x),St)
            res=ht.>Ct
            Vtj=zeros(J)
            for k=1:length(res)
                if res[k]==true
                    Vtj[k]=ht[k]
                else
                    Vtj[k]=Ct[k]
                end
            end
            sum_VT=sum(Vtj)/length(St)

            Q[t,i]=Q[t-1,i]/df+(Vt-sum_VT)
            U[t,i]=max(U[t-1,i]/df,(h[t,i]-Q[t,i]))

            if t==M+1
                U[t,i]=max(U[t-1,i]/df,mean(ht)-Q[t,i])
            end
        end
    end
    k=df^M
    U0=sum(U[M,:])/I2*(df^M)

end

#产生标的仿真路径-单路径
function generate_nest_mc(r::Float64,σ::Float64,dt::Float64,st::Float64,J::Int)

    p1=(r-0.5*(σ^2))*dt
    p2=rand(Normal(0,1),J)
    p3=p1 .+σ*√(dt)*p2
    S=st*(p3.|>exp)

    return S

end

#产生标的仿真路径-多路径
function genrate_path(S0::Float64,r::Float64,σ::Float64,dt::Float64,M::Int,I::Int)
    p1=(r-0.5*(σ^2))*dt
    S=zeros(M+1,I)
    S[1,:]=map(x->x=S0,S[1,:])
    @simd for t=2:M+1
        p2=rand(Normal(0,1),I)
        p3=p1 .+σ*√(dt)*p2
        S[t,:]=S[t-1,:].*(p3.|>exp)
    end

    return S
end

S0=36.0
K=40.0
T=1.0
r=0.06
σ=0.2
I1=16384
I2=1024

M=10
J=50
reg=5

res=LSM_DUAL(S0,K,T,r,σ,J,I1,I2,M,reg)
println(res)


using Distributions
using SpecialFunctions

export impliedVolatilityLiSORTS

function impliedVolatilitySqrtTimeRationalGuess(x::T, c::T) where {T}
    m00 = -0.00006103098165
    n00 = 1.0
    m01 = 5.33967643357688
    n01 = 22.96302109010794
    m10 = -0.40661990365427
    n10 = -0.48466536361620
    m02 = 3.25023425332360
    n02 = -0.77268824532468
    m11 = -36.19405221599028
    n11 = -1.34102279982050
    m20 = 0.08975394404851
    n20 = 0.43027619553168
    m03 = 83.84593224417796
    n03 = -5.70531500645109
    m12 = 41.21772632732834
    n12 = 2.45782574294244
    m21 = 3.83815885394565
    n21 = -0.04763802358853
    m30 = -0.21619763215668
    n30 = -0.03326944290044
    num = m00 + m10 * x + m20 * x^2 + m30 * x^3 + c * (m01 + m11 * x + m21 * x^2) + c^2 * (m02 + m12 * x) + c^3 * m03
    den = n00 + n10 * x + n20 * x^2 + n30 * x^3 + c * (n01 + n11 * x + n21 * x^2) + c^2 * (n02 + n12 * x) + c^3 * n03
    return num / den
end

function impliedVolatilityLiRationalGuess(
    price::T,
	isCall::Bool,
	strike::T,
    forward::T,
    timeToExpiry::T,
    df::T
) where {T}
    c = price / forward / df
    ex = forward / strike
    x0 = log(ex)
    x = x0
    if !isCall
        if x0 <= 0
            c = c + 1 - 1.0 / ex # put call parity
        else
			#duality + put call parity
            c = ex * c
            x = -x0
            ex = 1.0 / ex
        end
    else
        if x0 > 0
			# use c(-x0, v)
            c = (forward * (c - 1) + strike) / strike #in out duality
            x = -x0
            ex = 1.0 / ex
        end
    end
    if c >= 1 / ex || c >= 1
        error("price ", c, " higher than intrinsic value ", 1 / ex)
    elseif c <= 0
        error("price ", c, " negative")
    end
    sqrtte = sqrt(timeToExpiry)
    return impliedVolatilitySqrtTimeRationalGuess(x, c) / sqrtte
end

"Implied volatility of a vanilla option under the Black-Scholes model via Li SOR-TS algorithm"
function impliedVolatilityLiSORTS(
    price::T,
    isCall::Bool,
    strike::T,
    forward::T,
    timeToExpiry::T,
    df::T,
    impliedVolGuess::T,
    tolerance::T,
    maxIterations::Int
) where {T}
    c = price / forward / df
    ex = forward / strike
    x0 = log(ex)
    x = x0
    if !isCall
        if x0 <= 0
            c = c + 1 - 1.0 / ex # put call parity
        else
			#duality + put call parity
            c = ex * c
            x = -x0
            ex = 1.0 / ex
        end
    else
        if x0 > 0 			# use c(-x0, v)
            c = (forward * c - forward + strike) / strike # in out duality
            x = -x0
            ex = 1.0 / ex
        end
    end
    if c > 1 / ex
        return 0
    end
    sqrtte = sqrt(timeToExpiry)
    v0 = 0.0
    if impliedVolGuess == 0.0
        v0 = impliedVolatilitySqrtTimeRationalGuess(x, c)
    else
        v0 = impliedVolGuess * sqrtte
    end
    return computeLiSORTS(c, x, ex, v0, sqrtte, tolerance, maxIterations)
end

const OneOverSqrtTwo = sqrt(0.5)

function computeLiSORTS(c::T, x::T, ex::T, v0::T, sqrttte::T, tolerance::T, maxIterations::Int) where {T}
    v = v0
    if maxIterations <= 0
        return v
    end
    exinvsqrt = 1.0 / sqrt(ex)
    h = x / v
    t = v / 2

    Np = erfcx(-OneOverSqrtTwo * (h + t))
    Nm = erfcx(-OneOverSqrtTwo * (h - t))
    norm = exinvsqrt * 0.5 * exp(-0.5 * (h^2 + t^2))
    cEstimate = norm * (Np - Nm)
    if cEstimate < 0
        cEstimate = 0.0
    end

    omega = 1.0
    iterations = 0
    normal = Normal()
    while abs(c - cEstimate) > tolerance && iterations < maxIterations
        phi = (v^2 + 2 * x) / (v^2 - 2 * x)
        alpha = (1 + omega) / (1 + phi)
        F = c + norm * (Nm + omega * Np)
        Nm1 = quantile(normal, F / (1 + omega))
        G = Nm1 + sqrt(Nm1^2 - 2 * x)
        v = alpha * G + (1 - alpha) * v
        h = x / v
        t = v / 2
        Np = erfcx(-OneOverSqrtTwo * (h + t))
        Nm = erfcx(-OneOverSqrtTwo * (h - t))
        norm = exinvsqrt * 0.5 * exp(-0.5 * (h^2 + t^2))
        cEstimate = norm * (Np - Nm)
        if cEstimate < 0
            cEstimate = 0
        end
        iterations += 1
    end
    sigma = v / sqrttte
    return sigma
end


## fit splines to angle data to drive joints
using Dierckx, Interpolations

# fit quintic splines to data
joint_data = readdlm("data\\matchingData.csv", ',', Float64, skipstart=2)
hat_data = readdlm("data\\HAT.csv", ',', Float64, skipstart=2)
_time = joint_data[:, 1]
rh_spl = Spline1D(_time, joint_data[:, 4], k=5, s=0.0)
lh_spl = Spline1D(_time, joint_data[:, 5], k=5, s=0.0)
rk_spl = Spline1D(_time, joint_data[:, 6], k=5, s=0.0)
lk_spl = Spline1D(_time, joint_data[:, 7], k=5, s=0.0)
ra_spl = Spline1D(_time, joint_data[:, 8], k=5, s=0.0)
la_spl = Spline1D(_time, joint_data[:, 9], k=5, s=0.0)
rmtp_spl = Spline1D(_time, joint_data[:, 10], k=5, s=0.0)
lmtp_spl = Spline1D(_time, joint_data[:, 11], k=5, s=0.0)
hat_spl = Spline1D(hat_data[:, 1], hat_data[:, 2], k=5, s=0.0)

# get data for 1st and 2nd derivatives then interpolate
function fit_spl(time, spl; d2r=true)
    # evaluate 0th, 1st and 2nd derivative over time
    if d2r
        data0 = t -> evaluate(spl, t) .|> deg2rad
        data1 = t -> derivative(spl, t, 1) .|> deg2rad
        data2 = t -> derivative(spl, t, 2) .|> deg2rad
    else
        data0 = t -> evaluate(spl, t)
        data1 = t -> derivative(spl, t, 1)
        data2 = t -> derivative(spl, t, 2)
    end

    return data0, data1, data2

    # interpolate over data
    spl0 = CubicSplineInterpolation(time, data0, extrapolation_bc=Flat())
    spl1 = CubicSplineInterpolation(time, data1, extrapolation_bc=Flat())
    spl2 = CubicSplineInterpolation(time, data2, extrapolation_bc=Flat())

    return spl0, spl1, spl2
end

# generate functions
new_time = _time[1]:0.001:_time[end]
_rh, _rhp, _rhpp = fit_spl(new_time, rh_spl)
_lh, _lhp, _lhpp = fit_spl(new_time, lh_spl)
_rk, _rkp, _rkpp = fit_spl(new_time, rk_spl)
_lk, _lkp, _lkpp = fit_spl(new_time, lk_spl)
_ra, _rap, _rapp = fit_spl(new_time, ra_spl)
_la, _lap, _lapp = fit_spl(new_time, la_spl)
_rmtp, _rmtpp, _rmtppp = fit_spl(new_time, rmtp_spl)
_lmtp, _lmtpp, _lmtppp = fit_spl(new_time, lmtp_spl)
_hato, _hatop, _hatopp = fit_spl(new_time, hat_spl, d2r=false)


module Regression

using LinearAlgebra
using Statistics
export LinearRegression, Lasso, Ridge, make_design_matrix

include("utils.jl")
include("Base.jl")
include("BFM.jl")
include("LR.jl")
include("RR.jl")

end


using CurveFit
using Test
using LinearAlgebra
using Polynomials

lspf(start, stop, len) = range(start, stop=stop, length=len)

@testset "CurveFit.jl" begin

# Testing linear fit
x = lspf(1,10,10);
fun0(x) = 1.0 + 2.0*x
y = fun0.(x)
f = linear_fit(x, y)

@test f[1] ≈ 1.0 atol=1.0e-7
@test f[2] ≈ 2.0 atol=1.0e-7

f = curve_fit(LinearFit, x, y)
@test f(1.5) ≈ fun0(1.5) atol=1.0e-7

# power
fun1(x) = 1.0 + 2.0*log(x)
y = fun1.(x)
f = log_fit(x, y)
@test f[1] ≈ 1.0 atol=1.0e-7
@test f[2] ≈ 2.0 atol=1.0e-7
f = curve_fit(LogFit, x, y)
@test f(1.5) ≈ fun1(1.5) atol=1.0e-7

# log
fun2(x) = 2.0*x.^0.8
y = fun2(x)
f = power_fit(x, y)
@test f[1] ≈ 2.0 atol=1.0e-7
@test f[2] ≈ 0.8 atol=1.0e-7
f = curve_fit(PowerFit, x, y)
@test f(1.5) ≈ fun2(1.5) atol=1.0e-7

# Exp
fun3(x) = 2.0*exp.(0.8*x)
y = fun3(x)
f = exp_fit(x, y)
@test f[1] ≈ 2.0 atol=1.0e-7
@test f[2] ≈ 0.8 atol=1.0e-7
f = curve_fit(ExpFit, x, y)
@test f(1.5) ≈ fun3(1.5) atol=1.0e-7

# Poly
fun4(x) = 1.0 + 2.0*x + 3.0*x^2 + 0.5*x^3
y = fun4.(x)
f = poly_fit(x, y, 4)
@test f[1] ≈ 1.0 atol=1.0e-7
@test f[2] ≈ 2.0 atol=1.0e-7
@test f[3] ≈ 3.0 atol=1.0e-7
@test f[4] ≈ 0.5 atol=1.0e-7
@test f[5] ≈ 0.0 atol=1.0e-7

f = curve_fit(Polynomial, x, y, 4)
@test f(1.5) ≈ fun4(1.5) atol=1.0e-7

# Polynomials with large numbers
coefs = [80.0, -5e-18, -7e-20, -1e-36]
P = Polynomial(coefs)
x1 = 1e10 * (0:0.1:5)
y1 = P.(x1)
P2 = curve_fit(Polynomial, x1, y1, 3)
c = coeffs(P2)
@test coefs[1] ≈ c[1] rtol=1e-5
@test coefs[2] ≈ c[2] rtol=1e-5
@test coefs[3] ≈ c[3] rtol=1e-5
@test coefs[4] ≈ c[4] rtol=1e-5

# King's law
U = [lspf(1, 20, 20);]
A = 5.0
B = 1.5
n = 0.5
E = sqrt.(A .+ B*U.^n)
fun5(E) = ((E.^2 - A)/B).^(1 ./ n)

f = linear_king_fit(E, U)
@test f[1] ≈ A atol=1.0e-7
@test f[2] ≈ B atol=1.0e-7
f= curve_fit(LinearKingFit, E, U)
@test f(3.0) ≈ fun5(3.0) atol=1.0e-7

# Modified King's law
n = 0.42

E = sqrt.(A .+ B.*U.^n)
fun6(E) = ((E^2 - A)/B)^(1 / n)

f = king_fit(E, U)
@test f[1] ≈ A atol=1.0e-7
@test f[2] ≈ B atol=1.0e-7
@test f[3] ≈ n atol=1.0e-7
f= curve_fit(KingFit, E, U)
@test f(3.0) ≈ fun6(3.0) atol=1.0e-5

# Linear Rational fit
r =  RationalPoly([1.0, 0.0, -2.0], [1.0, 2.0, 3.0])
y = r.(x)
f = linear_rational_fit(x, y, 2, 3)
@test f[1] ≈ 1.0 atol=1.0e-8
@test f[2] ≈ 0.0 atol=1.0e-8
@test f[3] ≈ -2.0 atol=1.0e-8
@test f[4] ≈ 2.0 atol=1.0e-8
@test f[5] ≈ 3.0 atol=1.0e-8
@test f[6] ≈ 0.0 atol=1.0e-8

# Nonlinear Rational fit
f = rational_fit(x, y, 2, 3)

@test f[1] ≈ 1.0 atol=1.0e-7
@test f[2] ≈ 0.0 atol=1.0e-7
@test f[3] ≈ -2.0 atol=1.0e-7
@test f[4] ≈ 2.0 atol=1.0e-7
@test f[5] ≈ 3.0 atol=1.0e-7
@test f[6] ≈ 0.0 atol=1.0e-7

f = curve_fit(RationalPoly, x, y, 2, 3)
@test f(1.5) ≈ r(1.5) atol=1.0e-8
@test f(4.5) ≈ r(4.5) atol=1.0e-8

    # Gauss-Newton curve fitting. Linear problem:
    x = 1.0:10.0
    a0 = [3.0, 2.0, 1.0]
    fun7(x, a) = a[1] + a[2]*x + a[3]*x^2
    y = fun7.(x, Ref(a0))

    function ∇fun7!(x, a, df)
        df[1] = 1.0
        df[2] = x
        df[3] = x^2
    end

    a = gauss_newton_fit(x, y, fun7, ∇fun7!, [0.5, 0.5, 0.5], 1e-8, 30)

    @test a[1] ≈ a0[1] rtol=1e-7
    @test a[2] ≈ a0[2] rtol=1e-7
    @test a[3] ≈ a0[3] rtol=1e-7

    # Gauss-Newton curve fitting. Nonlinear problem:
    x = 1.0:10.0
    fun8(x, a) = a[1] + a[2] * x^a[3]
    a0 = [3.0, 2.0, 0.7]
    y = fun8.(x, Ref(a0))
    function ∇fun8!(x, a, df)
        df[1] = 1.0
        df[2] = x^a[3]
        df[3] = a[2] * x^a[3] * log(x)
    end
    a = gauss_newton_fit(x, y, fun8, ∇fun8!, [0.5, 0.5, 0.5], 1e-8, 30)

    @test a[1] ≈ a0[1] rtol=1e-7
    @test a[2] ≈ a0[2] rtol=1e-7
    @test a[3] ≈ a0[3] rtol=1e-7

    # Gauss-Newton curve fitting (generic interface). Linear problem:
    x = 1.0:10.0
    a0 = [3.0, 2.0, 1.0]
    fun9(x, a) = a[1] + a[2]*x[1] + a[3]*x[1]^2 - x[2]
    P = length(x)
    X = zeros(P, 2)
    for i in 1:P
        X[i,1] = x[i]
        X[i,2] = a0[1] + a0[2]*x[i] + a0[3]*x[i]^2
    end

    function ∇fun9!(x, a, df)
        df[1] = 1.0
        df[2] = x[1]
        df[3] = x[1]^2
    end

    a = gauss_newton_generic_fit(X, fun9, ∇fun9!, [0.5, 0.5, 0.5], 1e-8, 30)

    @test a[1] ≈ a0[1] rtol=1e-7
    @test a[2] ≈ a0[2] rtol=1e-7
    @test a[3] ≈ a0[3] rtol=1e-7

    # Gauss-Newton curve fitting (generic interface). Nonlinear problem:
    x = 1.0:10.0
    funA(x, a) = a[1] + a[2] * x[1]^a[3] - x[2]
    a0 = [3.0, 2.0, 0.7]
    P = length(x)
    X = zeros(P, 2)
    for i in 1:P
        X[i,1] = x[i]
        X[i,2] = a0[1] + a0[2]*x[i]^a0[3]
    end
    function ∇funA!(x, a, df)
        df[1] = 1.0
        df[2] = x[1]^a[3]
        df[3] = a[2] * x[1]^a[3] * log(x[1])
    end
    a = gauss_newton_generic_fit(X, funA, ∇funA!, [0.5, 0.5, 0.5], 1e-8, 30)

    @test a[1] ≈ a0[1] rtol=1e-7
    @test a[2] ≈ a0[2] rtol=1e-7
    @test a[3] ≈ a0[3] rtol=1e-7

    # Gauss-Newton curve fitting (generic interface). Nonlinear problem:
    U = 0.5:0.5:10
    a0 = [2.0, 1.0, 0.35]
    E = @. sqrt(a0[1] + a0[2]*U^a0[3])

    X = hcat(E, U)
    funB(x, a) = a[1] + a[2]*x[2]^a[3] - x[1]^2
    function ∇funB!(x, a, df)
        df[1] = 1.0
        df[2] = x[2]^a[3]
        df[3] = a[2]*x[2]^a[3] * log(x[2])
    end

    a = gauss_newton_generic_fit(X, funB, ∇funB!, [0.5, 0.5, 0.5], 1e-8, 30)

    @test a[1] ≈ a0[1] rtol=1e-7
    @test a[2] ≈ a0[2] rtol=1e-7
    @test a[3] ≈ a0[3] rtol=1e-7

    # Secant method NLS curve fitting. Linear problem:
    x = 1.0:10.0
    a0 = [3.0, 2.0, 1.0]
    funC(x, a) = a[1] + a[2]*x + a[3]*x^2
    y = funC.(x, Ref(a0))

    a = CurveFit.secant_nls_fit(x, y, funC, [0.5, 0.5, 0.5], 1e-8, 30)

    @test a[1] ≈ a0[1] rtol=1e-7
    @test a[2] ≈ a0[2] rtol=1e-7
    @test a[3] ≈ a0[3] rtol=1e-7

    # Gauss-Newton curve fitting. Nonlinear problem:
    x = 1.0:10.0
    funD(x, a) = a[1] + a[2] * x^a[3]
    a0 = [3.0, 2.0, 0.7]
    y = funD.(x, Ref(a0))
    a = CurveFit.secant_nls_fit(x, y, funD, [0.5, 0.5, 0.5], 1e-8, 30)

    @test a[1] ≈ a0[1] rtol=1e-7
    @test a[2] ≈ a0[2] rtol=1e-7
    @test a[3] ≈ a0[3] rtol=1e-7

end


println("# DPM TESTS...")

# run demo code at first
println(" * try demo...")

using DataFrames

# number of samples per cluster
N = 200

# create clusters
Data = DataFrame( x = randn(N), y = randn(N), class = "cluster1" )

# append second cluster
append!(Data, DataFrame( x = randn(N) + 5, y = randn(N) + 5, class = "cluster2" ));

D = 2 # 2 dimensional data
N = 400 # number of data points

# data matrix
X = zeros(N, D)

X[:,1] = convert(Array, Data[:x])
X[:,2] = convert(Array, Data[:y])

# init base distribution parameters
mu0 = mean(X, 1)
kappa0 = 9.0
nu0 = 5.0
Sigma0 = eye(D) * 10

# base distribution and concentration parameter (Gaussian with Normal Inverse Wishart Prior)
H = WishartGaussian(vec(mu0), kappa0, nu0, Sigma0)

# train Dirichlet Process Mixture Model
models = train(DPM(H), Gibbs(), RandomInitialisation(k = 10), X);
train(DPM(H), Gibbs(maxiter = 1000), KMeansInitialisation(), X)

# use $\alpha = 10.0$
train(DPM(H, α = 10.0), Gibbs(), KMeansInitialisation(), X)

# use Random Initialisation
train(DPM(H), Gibbs(), RandomInitialisation(), X)

# use 50 iterations and 10 burnin runs
train(DPM(H), Gibbs(maxiter = 50, burnin = 10),
	KMeansInitialisation(), X)

# access the log likelihood values of all iteration
LLH = map(model -> model.energy, models)

# access the weights of all iterations
W = map(model -> model.weights, models)

# access the distributions of all iterations
D = map(model -> model.distributions, models)

println("# FINISHED")


export SQG

"""
    SQG

Data structure to store arrays and parameters for the
Surface-Quasi-Geostrophic model

- rho : Background density
- g : Gravity
- buoyancy_freq_n : Background stratification
- odg_b : Amplitude of the buoyancy
"""
mutable struct SQG

    grid :: Grid
    rho :: Float64
    g :: Float64
    buoyancy_freq_n :: Float64
    odg_b :: Float64
    hv_order :: Int
    hv_val :: Float64

    b̂ :: Array{ComplexF64, 2}
    ψ̂ :: Array{ComplexF64, 2}

    u_x :: Array{Float64, 2}
    u_y :: Array{Float64, 2}
    û_x :: Array{ComplexF64, 2}
    û_y :: Array{ComplexF64, 2}

    a :: Array{Float64, 2}
    â :: Array{ComplexF64, 2}

    function SQG( grid, f0 )

        rho=1e3
        g = 9.81
        buoyancy_freq_n = 3 * f0
        odg_b = 1e-3
        hv_order = 8
        hv_val = 0

        nx, ny = grid.nx, grid.ny

        b̂ = zeros(ComplexF64, nx÷2+1, ny)
        ψ̂ = zeros(ComplexF64, nx÷2+1, ny)

        u_x = zeros(nx, ny)
        u_y = zeros(nx, ny)
        û_x = zeros(ComplexF64, nx÷2+1, ny)
        û_y = zeros(ComplexF64, nx÷2+1, ny)

        a = zeros(Float64, nx, ny)
        â = zeros(ComplexF64, nx÷2+1, ny)

        new( grid, rho, g, buoyancy_freq_n, odg_b, hv_order, hv_val,
             b̂, ψ̂, u_x, u_y, û_x, û_y, a, â)

    end

end

export update_velocities!

"""
    update_velocities!(u_x, u_y, model, b̂)

Compute the streamfunction and the velocity from the
buoyancy according to the Surface-Quasi-Geostrophic model.

"""
function update_velocities!( sqg :: SQG )

    nx, ny = sqg.grid.nx, sqg.grid.ny
    sqg.ψ̂ .= sqg.b̂ .* sqg.grid.on_k ./ sqg.buoyancy_freq_n

    sqg.û_x .= - 1im .* sqg.grid.ky .* sqg.ψ̂
    sqg.û_y	.=  1im .* sqg.grid.kx .* sqg.ψ̂

    sqg.u_x .= irfft(sqg.û_x, nx)
    sqg.u_y .= irfft(sqg.û_y, nx)

end

export update_advection_term!

"""
    update_advection_term(model)

Compute the Fourier transform of the partial derivative
of the advection term plus the hyperviscosity term

"""
function update_advection_term!(sqg)

    nx = sqg.grid.nx

    update_velocities!( sqg )

    # Advection term

    sqg.a .=  sqg.u_x .* irfft(-1im .* sqg.grid.kx .* sqg.b̂, nx)
    sqg.a .+= sqg.u_y .* irfft(-1im .* sqg.grid.ky .* sqg.b̂, nx)

    sqg.â .= rfft(sqg.a)

    # Summing Hyperviscosity term

    sqg.â .-= sqg.hv_val .* (sqg.grid.k.^sqg.hv_order) .* sqg.b̂

end

export TimeSolver

struct TimeSolver

    b̂ :: Array{ComplexF64, 2}
    db̂ :: Array{ComplexF64, 2}

    function TimeSolver( nx, ny)

        b̂ = zeros(ComplexF64, nx÷2+1, ny)
        db̂ = zeros(ComplexF64, nx÷2+1, ny)

        new( b̂, db̂ )

    end

end

export step!

"""
    step!(model, time_solver, dt)

Time integration of the advection equation of b
by 4th order Runge-Kutta method
"""
function step!( sqg :: SQG, rk4 :: TimeSolver, dt )

    rk4.b̂ .= sqg.b̂

    update_advection_term!(sqg)

    sqg.b̂ .= rk4.b̂ .+ sqg.â .* dt / 2
    rk4.db̂ .= sqg.â ./ 2

    update_advection_term!(sqg)

    sqg.b̂ .= rk4.b̂ .+ sqg.â .* dt / 2
    rk4.db̂ .+= sqg.â

    update_advection_term!(sqg)

    sqg.b̂ .= rk4.b̂ .+ sqg.â .* dt
    rk4.db̂ .+= sqg.â

    update_advection_term!(sqg)

    rk4.db̂ .+= sqg.â ./ 2

    sqg.b̂ .= rk4.b̂ .+ dt / 3 .* rk4.db̂

end


using myJuliaUtils
using Test
using LinearAlgebra
using Random

# testing some helper function

# one matrix case (qq is a vector)
n = round(Int, 2 + rand() * 100)
qq = genPDmat(n)
qq = vech(qq)

A, sigma = myJuliaUtils.decVCV(qq)

@test vec2sym(qq) ≈ vec2ltriW1(A) * Diagonal(sigma).^2 * vec2ltriW1(A)'

# many matrix case (qq is a matrix)

n = 5
T = 100
c = round(Int, n * (n + 1) / 2)
qq = Array{Float64}(undef, c, T)

for i = 1:T
    q = genPDmat(n)
    qq[:,i] = vech(q)
end

A, sigma = myJuliaUtils.decVCV(qq)

for i = 1:T
    @test vec2sym(qq[:,i]) ≈ vec2ltriW1(A[:,i]) * Diagonal(sigma[:,i]).^2 * vec2ltriW1(A[:,i])'
end

A, sigma = myJuliaUtils.decVCV(qq,true)

for i = 1:T
    @test vec2sym(qq[:,i]) != vec2ltriW1(A[:,i]) * Diagonal(sigma[:,i]).^2 * vec2ltriW1(A[:,i])'
end


import Base.Broadcast.broadcasted
export QolElemAtom, qol_elementwise, square, sumsquares, invpos, /
export sign, monotonicity, curvature, conic_form!

struct QolElemAtom <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{AbstractExpr, AbstractExpr}
    size::Tuple{Int, Int}

    function QolElemAtom(x::AbstractExpr, y::AbstractExpr)
        if x.size != y.size
            error("elementwise quad over lin must take two arguments of the same size")
        end
        children = (x, y)
        return new(:qol_elem, hash(children), children, x.size)
    end
end

function sign(q::QolElemAtom)
    return Positive()
end

function monotonicity(q::QolElemAtom)
    return (sign(q.children[1]) * Nondecreasing(), Nonincreasing())
end

function curvature(q::QolElemAtom)
    return ConvexVexity()
end

function evaluate(q::QolElemAtom)
    return (evaluate(q.children[1]).^2) ./ evaluate(q.children[2])
end

function conic_form!(q::QolElemAtom, unique_conic_forms::UniqueConicForms)
    if !has_conic_form(unique_conic_forms, q)
        sz = q.children[1].size
        t = Variable(sz[1], sz[2])
        qol_objective = conic_form!(t, unique_conic_forms)
        x, y = q.children
        conic_form!(SOCElemConstraint(y + t, y - t, 2 * x), unique_conic_forms)
        # add implicit constraint y >= 0
        conic_form!(y >= 0, unique_conic_forms)
        cache_conic_form!(unique_conic_forms, q, qol_objective)
    end
    return get_conic_form(unique_conic_forms, q)
end

qol_elementwise(x::AbstractExpr, y::AbstractExpr) = QolElemAtom(x, y)

broadcasted(::typeof(^),x::AbstractExpr,k::Int) = k==2 ? QolElemAtom(x, Constant(ones(x.size[1], x.size[2]))) : error("raising variables to powers other than 2 is not implemented")

invpos(x::AbstractExpr) = QolElemAtom(Constant(ones(x.size[1], x.size[2])), x)
broadcasted(::typeof(/), x::Value, y::AbstractExpr) = DotMultiplyAtom(Constant(x), invpos(y))
/(x::Value, y::AbstractExpr) = size(y) == (1,1) ? MultiplyAtom(Constant(x), invpos(y)) : error("cannot divide by a variable of size $(size(y))")
sumsquares(x::AbstractExpr) = square(norm2(x))

function square(x::AbstractExpr)
    if sign(x) == ComplexSign()
        error("Square of complex number is not DCP. Did you mean square_modulus?")
    else
        QolElemAtom(x, Constant(ones(x.size[1], x.size[2])))
    end
end


function getfeatures_series(serie, prefix=""; sweights=ones(length(serie)), hasisless=true)
    features = Dict()
    isempty(serie) && return features
    features["$(prefix)mean"] = mean(serie, weights(sweights))
    !hasisless && return features
    features["$(prefix)var"] = var(serie, weights(sweights))
    features["$(prefix)median"] = median(serie)
    features["$(prefix)max"] = maximum(serie)
    features["$(prefix)min"] = minimum(serie)
    return features
end

"""
    getfeatures_graph(graph::AbstractGraph, vweights=ones(nv(graph)))

Get the graph features `graph`. Some features (denoted (w)) can be weighted per vertex using `vweights`.
Those features are:
    * Number of vertex
    * Number of edges
    * Degree mean (w)
    * Degree variation (w)
    * Degree max
    * Degree min
    * Density
"""
function getfeatures_graph(graph::AbstractGraph, vweights=ones(nv(graph)))
    graph_features = Dict()
    graph_features["nv"] = nv(graph)
    graph_features["ne"] = ne(graph)
    graph_features["degree_mean"] = mean(degree(graph), weights(vweights))
    graph_features["degree_var"] = var(degree(graph), weights(vweights))
    graph_features["degree_max"] = Δ(graph)
    graph_features["degree_min"] = δ(graph)
    graph_features["density"] = density(graph)
    graph_features["diameter"] = diameter(graph)
    graph_features["radius"] = radius(graph)
    graph_features["global_clustering_coefficient"] = global_clustering_coefficient(graph)
    #graph_features["assortativity"] = assortativity(graph)
    #graph_features["eccentricity"] = eccentricity(graph)
    return graph_features
end

"""
    getfeatures_cliques(graph::AbstractGraph, cliques, cweights=ones(length(cliques)))

Get the cliques features of `graph`. Some features (denoted (w)) can be weighted per vertex using `vweights`.
Those features are:
    * Number of cliques
    * Clique size mean (w)
    * Clique size variation (w)
    * Clique size max
    * Clique size min
"""
function getfeatures_cliques(graph::AbstractGraph, cliques, cweights=ones(length(cliques)))
    cliques_features = Dict()
    cliques_features["nb"] = length(cliques)
    cliques_size = [length(clique) for clique in cliques]
    merge!(cliques_features, getfeatures_series(cliques_size, "size_"))
    return cliques_features
end

"""
    getfeatures_kernel(graph::AbstractGraph, vweights=ones(nv(graph)))

Get the kernel features of `graph`. The kernel being a subgraph of `graph` those features corresponds to graph features.
See also [`get_graph_features`].
"""
function getfeatures_kernel(graph::AbstractGraph, k=1)
    kernel_graph = copy(graph)
    vs = Set(k_core(kernel_graph, k))
    for v in vertices(kernel_graph)
        if !(v in vs)
            rem_vertex!(kernel_graph, v)
        end
    end
    return getfeatures_graph(kernel_graph)
end

"""
    getfeatures_OPF(path_matpower)

Get the OPF features from a matpower file.
Those features are:
    * The generators costs
    * The loads
    * The shunts
    * The voltage bounds
    * The realpower bounds
    * The max current
    * The generators density
"""

function getfeatures_OPF(path_matpower)
    costs_generators, loads, shunts, bounds_voltage, bounds_realpower, bounds_imagpower, max_current, generator_density = get_data_OPF(path_matpower)
    costs_generators = [[parse(Int, key[5:end]), costs_generators[key]...] for key in keys(costs_generators)]
    loads = sort([[parse(Int, key[5:end]), loads[key]] for key in keys(loads)], by=x->real(x[1]))
    loads = map(x->x[2], loads)
    shunts = sort([[parse(Int, key[5:end]), shunts[key]] for key in keys(shunts)], by=x->real(x[1]))
    shunts = map(x->x[2], shunts)
    bounds_voltage = sort([[parse(Int, key[5:end]), bounds_voltage[key]...] for key in keys(bounds_voltage)])
    bounds_voltage = map(x->x[2:end], bounds_voltage)
    bounds_realpower = sort([[parse(Int, key[5:end]), bounds_realpower[key]...] for key in keys(bounds_realpower)])
    bounds_realpower = map(x->x[2:end], bounds_realpower)
    bounds_imagpower = sort([[parse(Int, key[5:end]), bounds_imagpower[key]...] for key in keys(bounds_imagpower)])
    bounds_imagpower = map(x->x[2:end], bounds_imagpower)
    max_current = sort([[parse(Int, key[1][5:end]), parse(Int, key[1][5:end]), val] for (key, val) in max_current], by=x->x[begin:2])
    features = Dict("costs_generators_linear" => getfeatures_series(map(x -> x[1], costs_generators)),
                    "costs_generators_constant" => getfeatures_series(map(x -> x[2], costs_generators)),
                    "loads_real" => getfeatures_series(map(x->real(x), loads)),
                    "loads_imag" => getfeatures_series(map(x->imag(x), loads)),
                    "shunts_real" => getfeatures_series(map(x->real(x), shunts)),
                    "shunts_imag" => getfeatures_series(map(x->imag(x), shunts)),
                    "bounds_voltage_low" => !isempty(bounds_voltage) ? getfeatures_series(map(x -> x[1], bounds_voltage)) : Dict(),
                    "bounds_voltage_up" => !isempty(bounds_voltage) ? getfeatures_series(map(x -> x[2], bounds_voltage)) : Dict(),
                    "bounds_realpower_low" => !isempty(bounds_realpower) ? getfeatures_series(map(x -> x[1], bounds_realpower)) : Dict(),
                    "bounds_realpower_up" => !isempty(bounds_realpower) ? getfeatures_series(map(x -> x[2], bounds_realpower)) : Dict(),
                    "bounds_imagpower_low" => !isempty(bounds_imagpower) ? getfeatures_series(map(x -> x[1], bounds_imagpower)) : Dict(),
                    "bounds_imagpower_up" => !isempty(bounds_imagpower) ? getfeatures_series(map(x -> x[2], bounds_imagpower)) : Dict(),
                    "max_current" => !isempty(max_current) ? getfeatures_series(map(x -> x[end], max_current)) : Dict(),
                    "generator_density" => generator_density)
    return features
end

"""
    get_data_OPF(instance_path::String)

Retrieve the OPF related data from a matpower file.
"""
function get_data_OPF(instance_path::String)
    data = load_matpower(instance_path)
    costs_generators = Dict{String, Tuple{Float64, Float64}}() # bus g => (cg, kg)
    loads = Dict{String, Complex}() # bus => Sl= Pl + im*Ql
    shunts = Dict{String, Complex}() # bus n => g - im *b
    bounds_voltage = Dict{String, Tuple{Float64, Float64}}() # bus => (vmin, vmax)
    bounds_realpower = Dict{String, Tuple{Float64, Float64}}() # bus => (Pmin, Pmax)
    bounds_imagpower = Dict{String, Tuple{Float64, Float64}}() # bus => (Qmin, Qmax)
    max_current = Dict{Tuple{String, String}, Float64}() # line=(bus, bus) => imax
    bus_id_line=SortedDict{Int, Int}()
    bus_id_name=SortedDict{Int, String}()
    checkfor(data, 2, "mpc.baseMVA")
    baseMVA = data[2, 3]
    ## Building bus load and shunt information
    i_debut, i_fin = find_numarray(1, data)
    bustype = data[i_debut:i_fin, 2]                # weather bus is 1:"PQ" (generators are not accounted for) or 2:"PV"
    checkfor(data, i_debut-1, "mpc.bus")
    for i in i_debut:i_fin
        id = i-i_debut+1
        busname = "BUS_$id"
        bounds_voltage[busname] = (data[i, 13], data[i, 12])
        loads[busname] = data[i, 3] + im * data[i, 4]
        if data[i, 5] == data[i, 6] == 0
            #no shunts
        else
            shunts[busname] = data[i, 5] - im * data[i, 6]
        end
        bus_id_line[data[i, 1]] = id
        bus_id_name[data[i, 1]] = busname
    end
    ## Adding bus generator information
    gen2bus = SortedDict{Int, Int}()
    line2busgen = SortedDict{Int, Tuple{Int, Int}}()
    i_debut, i_fin = find_numarray(i_fin + 1, data)
    checkfor(data, i_debut - 1, "mpc.gen")
    genind, prevgen = 0, 0
    for i=i_debut:i_fin
        gen2bus[i - i_debut + 1] = bus_id_line[data[i, 1]]
        busname = bus_id_name[data[i, 1]]
        S_min = S_max = 0
        if data[i, 1] == prevgen
            genind += 1
            (Pmin, Pmax) = bounds_realpower[busname]
            (Qmin, Qmax) = bounds_imagpower[busname]
            S_min = Pmin + im * Qmin
            S_max = Pmax + im * Qmax
        else
            prevgen = data[i, 1]
            genind = 1
        end
        line2busgen[i - i_debut + 1] = (data[i, 1], genind)
        if data[i, 8] > 0 #generator is on
            S_min += data[i, 10] + im * data[i, 5] #Smin = Pmin + i Qmin
            S_max += data[i, 9] + im * data[i, 4] #Smax = Pmax + i Qmax
        end
        bounds_realpower[busname] = (real(S_min), real(S_max))
        bounds_imagpower[busname] = (imag(S_min), imag(S_max))
    end
    ## building link information
    i_debut, i_fin = find_numarray(i_fin + 1, data)
    checkfor(data, i_debut - 1, "mpc.branch")
    for i=i_debut:i_fin
        if data[i, 11] == 0
            #@warn("$link $(linkname.orig)⟶$(linkname.dest) breaker out of service !")
        else
            rs, xs, bc = data[i, 3:5]
            Smax = data[i, 6]
            τ, θ = data[i, 9:10]
            if Smax != 0
                max_current[(bus_id_name[data[i, 1]], bus_id_name[data[i, 2]])] = Smax / baseMVA
            end
        end
    end
    i_debut, i_fin = find_numarray(i_fin+1, data)
    if data[i_debut - 1, 1] == "mpc.areas"
        @warn("Not loading mpc.areas data.")
        i_debut = i_fin + 1
        i_debut, i_fin = find_numarray(i_fin + 1, data)
    end
    bus_line_id = SortedDict([val=>key for (key, val) in bus_id_line])
    ## Adding generator cost information
    checkfor(data, i_debut - 1, "mpc.gencost")
    genind, cur_genind = 0, data[i_debut, 1]
    for i = i_debut:i_fin
        buslineid, _ = line2busgen[i - i_debut+1]
        busid = bus_id_line[buslineid]
        busname = "BUS_$busid"
        cost_degree = data[i, 4]
        cost_coeffs = data[i, 5:(5+cost_degree-1)]
        if cost_degree == 2
            costs_generators[busname] = (cost_coeffs[1], cost_coeffs[2])
        elseif cost_degree == 3
            costs_generators[busname] = (cost_coeffs[2], cost_coeffs[3])
        end
    end
    return costs_generators, loads, shunts, bounds_voltage, bounds_realpower, bounds_imagpower, max_current, NaN
end

function get_nonzeros_factor(path::String)
    file = split(read(open(path, "r"), String), "nonzeros before factor")
    (length(file) < 2) && return -1
    file = split(file[2], ' ')
    return Int(parse(Float64, file[3]))
end


module ConeProgramDiff

using LinearAlgebra, SparseArrays                       #stdlib
using JuMP, MathOptInterface, MathOptSetDistances       #Opt utils
using IterativeSolvers, Roots                           #Eq solvers
using SCS, ECOS, Hypatia                                #Opt solvers
using Random

const MOI = MathOptInterface
const MOSD = MathOptSetDistances

include("utils.jl")
include("cones.jl")
include("cone_solve.jl")
include("generate_program.jl")

export solve_and_diff
export project_onto_cone, d_project_onto_cone
export random_cone_program, l1_minimization_program

end

# TODO: Check out https://github.com/Jutho/LinearMaps.jl for projections


## Caculate Trace of Hat matrix in LMM (only for PLSGeneral type in MixedModels)

function Λ0(m::LinearMixedModel)
	vv = SparseMatrixCSC{Float64,Int}[]
	for i in 1:length(m.λ)
		nl = length(m.facs[i].pool)
		ll = m.λ[i]
		#p = size(ll,1)
		p = size(ll)[1]
		if p == 1
			isa(ll,MixedModels.PDScalF) || error("1×1 λ section should be a PDScalF type")
			push!(vv,ll.s .* speye(nl))
		else
			isa(ll, MixedModels.PDDiagF) ? sfll = sparse(full(ll.d)): sfll = sparse(full(ll.ch.UL))
			push!(vv,blkdiag([sfll for _ in 1:nl]...))
		end
	end
	Triangular(blkdiag(vv...),:L,false)
end

function traceHat(m::LinearMixedModel)
	X=m.X.m
	XtX = Symmetric(X'X,:L)
	Zt=MixedModels.zt(m)
	Λt=Λ0(m).data

	Ztc = CholmodSparse(Λt*Zt,0)
	L = cholfact(Ztc,1.,true)

	perm=L.Perm .+ one(eltype(L.Perm))
	RX=cholfact(MixedModels.symcontents(XtX),:L)
	RZX = Λt*Zt*X

	copy!(RZX,solve(L, RZX[perm,:], CHOLMOD_L))

	LAPACK.potrf!('L',BLAS.syrk!('L','T',-1.,RZX, 1.,
		copy!(RX.UL, MixedModels.symcontents(XtX))))

	Cl= solve(L,solve(L,Ztc,CHOLMOD_P),CHOLMOD_L)
	#Cl= solve(L,Ztc,CHOLMOD_L)
	Cr=Ac_ldiv_B(tril(RX.UL), X' - full(Ac_mul_B(Cl,RZX))' )

	rho=sumabs2(Cl.nzval)+sumabs2(vec(Cr))
	rho
end

function vinv(m::LinearMixedModel)

	Zt=MixedModels.zt(m)
	#Λ=SelectMixedModels.Λ0(m)

    L1 = cholfact(CholmodSparse(SelectMixedModels.Λ0(m).data,0),true)
    T=solve(L1,CholmodSparse(speye(size(L1)[1]),0), CHOLMOD_A)
	L2 = cholfact(CholmodSparse(sparse(T) + Zt *Zt'),true)
	S = solve(L2,CholmodSparse(speye(size(L2)[1]),0), CHOLMOD_A)
	return  speye(size(m)[1]) - Zt' * sparse(S) * Zt
end


using ProximalOperators
Random.seed!(2)

xsol1 = randn(100)

A = randn(50,100)
b = A*xsol1

S1 = IndAffine(A,b)
S2 = IndBox(0.0, Inf)

problem = Feasibility(S1, S2, 100)

sol, model = FirstOrderSolvers.solve!(problem, DR(eps=1e-8, verbose=0),checki=10)

@test sol.status == :Optimal
@test minimum(sol.x) > -1e-12
@test maximum(abs, A*sol.x - b) < 1e-12

sol, model = FirstOrderSolvers.solve!(problem, AP(eps=1e-8, verbose=0),checki=1)

@test sol.status == :Indeterminate

sol, model = FirstOrderSolvers.solve!(problem, GAP(eps=1e-8, verbose=0))

@test sol.status == :Indeterminate

sol, model = FirstOrderSolvers.solve!(problem, FISTA(eps=1e-8, verbose=1))

@test sol.status == :Indeterminate

@testset "Solver: $(typeof(alg))" for alg in
    [GAPP(eps=1e-8, verbose=0, proji=50),
    GAPA(eps=1e-8, verbose=0),
    LineSearchWrapper(GAP(eps=1e-8, verbose=0))]

    sol, model = FirstOrderSolvers.solve!(problem, alg)

    @test sol.status == :Optimal
    @test minimum(sol.x) > -1e-12
    @test maximum(abs, A*sol.x - b) < 1e-6
end


@userplot MtxExpsPlot

RecipesBase.@recipe function f(p::MtxExpsPlot; dt=0.1, nmax=50, lbmethod=:none)
    if length(p.args) != 1 || !(typeof(p.args[1]) <: PSAStruct)
        error("mtxexpsplot must be given a PSAStruct. Got $(typeof(p.args))")
    end
    ps_data::PSAStruct = p.args[1] # NB hint the type to the compiler
    A = ps_data.input_matrix

    ts, transient = get_mtxexpnorm(A, dt, nmax)
    alp = maximum(real.(ps_data.ps_dict[:ews]))

    #layout := @layout [a; b]
    #link := :x
    seriestype := :path
    label := ""

    @series begin
        #subplot := 1
        label := "Matrix norm"
        ts, transient
    end
    @series begin
        #subplot := 1
        label := "Spectral Radius"
        ts, exp.(alp*ts)
    end
    #=
    @series begin
        subplot := 2
        yscale := :log10
        ts, transient
    end
    @series begin
        subplot := 2
        ts, exp.(alp*ts)
    end
    =#
end


concept_dist_different(x) = abs(x[1] - x[2]) ≠ abs(x[3] - x[4])

"""
    dist_different
Local constraint ensuring that `concept(dist_different, x) = |x[1] - x[2]| ≠ |x[3] - x[4]|)`.
"""
const dist_different = Constraint(
    args_length = 4,
    concept = concept_dist_different,
    error = make_error(:dist_different),
)


using BayesianLinearRegression

import Random
import Printf: @printf

Random.seed!(2)
n = 20;
xx = range(0, 2π, length=n);
y = sin.(xx) .+ 0.1 .* randn(n);

function compareModels(p)
    X = chebyshev(xx,p)
    m = BayesianLinReg(X,y)
    fit!(m, callback=stopAtIteration(10))
    return logEvidence(m)
end

for (p,logEv) in enumerate(compareModels.(1:20))
    @printf "p = %2d: evidence = %8.2f\n" p exp(logEv)
end

X = chebyshev(xx,6);
@time m = BayesianLinReg(X,y) |> fit!;

priorScale(m)

noiseScale(m)

posteriorWeights(m)

effectiveNumParameters(m)

yhat = predict(m,m.X)


using Church

dp(concentration::Real, base_measure::Function) = begin
    sticks = Mem(i::Int -> beta(1., concentration))
    atoms  = Mem(i::Int -> base_measure())
    loop(i::Int) =
        @If(bernoulli(sticks[i]), atoms[i], loop(i+1))
    d = () -> loop(1)
end

dp_mixture(concentration::Real, base_measure::Function, parameter::Function) = begin
    dp_ = dp(concentration, base_measure)
    () -> parameter(dp_())
end

d = dp(1., normal)
ds = [d() for i = 1:20]

f = dp_mixture(1., normal, x -> normal(x, 0.1))
fs = [f() for i = 1:20]

n_data = 10
n_components = 20
indicies = [categorical(n_components) for i = 1:n_data]
components = Mem((i::Int) -> normal())
data = [components[indicies[i]] for i = 1:n_data]