module NEPTransformations

    # Transformations between types
    using SparseArrays

    include("nleigs_coefficients.jl")

    export taylor_expansion_pep
    export shift_and_scale
    export mobius_transform
    export CORKPencil
    export buildPencil
    export CorkLinearization
    export DefaultCorkLinearization
    export IarCorkLinearization
    export NleigsCorkLinearization
    export CORKPencilLR
    export lowRankCompress

    # We overload these
    import ..NEPCore.compute_Mder
    import ..NEPCore.compute_Mlincomb
    import ..NEPCore.compute_Mlincomb!
    import ..NEPCore.compute_MM
    import ..NEPCore.compute_resnorm
    import Base.size
    import SparseArrays.issparse

    """
        struct ShiftScaledNEP <: NEP
        ShiftScaleNEP(orgnep::NEP[,shift=0][,scale=1])

    Transforms a nep (orgnep) M(λ)v to a new nep T(λ)=M(scale*λ+shift). This can be used if the method does not have an easy implementation of shift and scaling. Usage of this transformation can slow down the algorithm.
    """
    struct ShiftScaledNEP <: NEP
        shift::Number
        scale::Number
        orgnep::NEP
        function ShiftScaledNEP(orgnep;shift=0,scale=1)
            return new(shift,scale,orgnep)
        end
    end
    """
        shift_and_scale(orgnep::NEP;shift=0,scale=1)

    Transforms the orgnep by defining a new NEP from the relation
    T(λ)=M(scale * λ+shift) where M is the orgnep. This function tries
     to preserve the NEP type, e.g., a shift_and_scale operation on
    an SPMF-object, return an SPMF object. If it cannot preserve
    the type, it will return a nep of the struct `ShiftScaledNEP`.

    # Example
    ```julia-repl
    julia> nep0=nep_gallery("pep0")
    julia> σ=3; α=10;
    julia> nep1=shift_and_scale(nep0,shift=σ,scale=α)
    julia> opnorm(compute_Mder(nep0,α*(4+4im)+σ)-compute_Mder(nep1,4+4im))
    8.875435870738592e-12
    ```
    """
    function shift_and_scale(orgnep::NEP;shift=0,scale=1)
        return ShiftScaledNEP(orgnep,shift=shift,scale=scale);
    end

    # Native implementation for PEPs by transformation of the coefficients. This is not necessarily fast.
    function shift_and_scale(orgnep::PEP;shift=0,scale=1)
        At=copy(shift*scale*orgnep.A); # Allocate correct type (not so fast)
        m=size(At,1)-1
        for j=0:m
            AA = zero(At[j+1])
            for i=j:m
                factor=((scale^j)*(shift^(i-j))*factorial(i)/factorial(i-j))/factorial(j)
                AA+=orgnep.A[i+1]*factor;

            end
            At[j+1]=AA;
        end
        return PEP(At)
    end

    # A shift and rescale of a DEP is DEP where the matrix coefficients and the delays are rescaled and another matrix coefficient with delay zero is added.
    function shift_and_scale(orgnep::DEP;shift=0,scale=1)
        J = Matrix{eltype(orgnep.A[1])}(I, size(orgnep.A[1]))
        return DEP([broadcast(*,orgnep.A,exp.(-orgnep.tauv*shift)/scale); [-shift/scale * J] ],[orgnep.tauv*scale; 0])
    end

    # Native implementation for SPMF. Only for generic SPMF (not AbstractSPMF)
    function shift_and_scale(orgnep::SPMF_NEP;shift=0,scale=1)
        orgfv=get_fv(orgnep);
        m=size(orgfv,1);
        fv = Vector{Function}(undef, m)
        # create anonymous functions corresponding to the
        # shift and scaled problem
        for i=1:m
            fv[i] = S -> orgfv[i](scale*S + shift*one(S))::((S isa Number) ? Number : Matrix)
        end
        # create a new SPMF with transformed functions
        return SPMF_NEP(get_Av(orgnep),fv;Schur_fact=orgnep.Schur_factorize_before);
    end

    # Size and issparse implemented by delegation
    function size(nep::ShiftScaledNEP,dim=-1)
        return size(nep.orgnep,dim)
    end
    function issparse(nep::ShiftScaledNEP)
        return issparse(nep.orgnep)
    end

    function compute_Mder(nep::ShiftScaledNEP,λ::Number,i::Integer=0)
        # Just scale the orgnep call. Chain rule for differentiation.
        return (nep.scale^i)*compute_Mder(nep.orgnep,nep.scale*λ+nep.shift,i);
    end

    function compute_MM(nep::ShiftScaledNEP,S,V)
        # Just call orgnep with a different S
        return compute_MM(nep.orgnep,S*nep.scale+eye(S)*nep.shift,V)
    end

    function compute_Mlincomb(nep::ShiftScaledNEP,λ::Number,V::AbstractVecOrMat,
                              a::Vector=ones(size(V,2)))
        # Multiply the coefficient matrix V with a diagonal matrix
        # corresponding to the chain rule
        p=size(V,2);
        z=nep.scale.^Array{eltype(V),1}(0:p-1)
        W=V * diagm(0 => z) # not very fast
        return compute_Mlincomb(nep.orgnep,nep.scale*λ+nep.shift,W,a);
    end

    """
        mobius_transform(orgnep::NEP,[,a=1][,b=0][,c=0][,d=1])

    Transforms a nep (orgnep) ``M(λ)v`` to a new nep ``T(λ)=M((aλ+b)/(cλ+d))``.
    This function tries to preserve the type such that `T`
    and `M` are of the same NEP-type (see `shift_and_scale()`).
    If it cannot be preserved it will return a `MobiusTransformedNEP`.
    It is in general advised to try to preserve the type, and
    the use of `MobiusTransformedNEP` can
    considerably slow down NEP-access.

    # Example

    ```julia-repl
    julia> nep0=nep_gallery("pep0")
    julia> a=1; b=3; c=4; d=5;
    julia> nep1=mobius_transform(nep0,a=a,b=b,c=c,d=d);
    julia> s=3;
    julia> opnorm(compute_Mder(nep0,(a*s+b)/(c*s+d))-compute_Mder(nep1,s))
    0.0
    ```
    """
    function mobius_transform(orgnep::NEP;a=1,b=0,c=0,d=1)
        return MobiusTransformedNEP(orgnep;a=a,b=b,c=c,d=d)
    end

    function mobius_transform(orgnep::SPMF_NEP;a=1,b=0,c=0,d=1)
        orgfv=get_fv(orgnep);
        m=size(orgfv,1);
        fv = Vector{Function}(undef, m)
        # create anonymous functions corresponding to the
        # möbius transformed problem
        for i=1:m
            fv[i] = S -> orgfv[i]((a*S + b*one(S)) / (c*S + d*one(S)))::((S isa Number) ? Number : Matrix)
        end
        # create a new SPMF with transformed functions
        return SPMF_NEP(get_Av(orgnep),fv;Schur_fact=orgnep.Schur_factorize_before);
    end

    """
    Hidden type representing a transformed NEP
    """
    struct MobiusTransformedNEP <: NEP
        a::Number
        b::Number
        c::Number
        d::Number
        orgnep::NEP
        function MobiusTransformedNEP(orgnep;a=1,b=0,c=0,d=1)
            return new(a,b,c,d,orgnep)
        end
    end

    # Size and issparse implemented by delegation
    function size(nep::MobiusTransformedNEP,dim=-1)
        return size(nep.orgnep,dim)
    end
    function issparse(nep::MobiusTransformedNEP)
        return issparse(nep.orgnep)
    end

    function compute_Mder(nep::MobiusTransformedNEP,λ::Number,i::Integer=0)
        # I did not found a better way
        return compute_Mder_from_MM(nep,λ,i);
    end

    function compute_MM(nep::MobiusTransformedNEP,S,V)
        # Just call orgnep with a different S
        return compute_MM(nep.orgnep, (nep.a*S + nep.b*I) / (nep.c*S + nep.d*I), V)
    end

    function compute_Mlincomb(nep::MobiusTransformedNEP,λ::Number,V::AbstractVecOrMat
                              ,a::Vector=ones(size(V,2)))
        # I did not found a better way
        return compute_Mlincomb_from_MM(nep,λ,V,a)
    end

    # consider renaming this function
    """
        taylor_expansion_pep(orgnep::NEP[,d=2])

    Compute the truncated (with `d` terms) Taylor series of the NEP.
    The output is a [`PEP`](@ref).
    """
    function taylor_expansion_pep(nep::NEP,d::Integer=2)
        A=Vector{Matrix}(undef,d+1)
        for i=0:d
            A[i+1]=compute_Mder(nep,0,i)/factorial(i);
        end
        return PEP(A)
    end

"""
    struct CORKPencil
    function CORKPencil(M,N,Av,Bv,Z);
    function CORKPencil(nep,is)

The struct `CORKPencil` represents a pencil with a particular structure,
as given in the reference. The data can either be constructed directly
via the first constructor, or from a NEP in the second constructor.
The second constructor takes a `NEP` and  `is` which specifies
a CORK-structure as well as an approximation method. This can be objects
of the type `IarCorkLinearization` or `NleigsCorkLinearization`,
which are the CORK-linearizations equivalent to (certain versions of)
[`iar`](@ref) and [`nleigs`](@ref).

See [`buildPencil`](@ref) how to build standard pencil.

# Example:

The following example constructs a  `CORKPencil` from
a general NEP and then computes approximations of NEPs
by the interpolation approach of `nleigs`.

```julia-repl
julia> using LinearAlgebra
julia> A=(1:4)*(1:4)'+I; B=diagm(1 => [1,2,3]); C=ones(4,4);
julia> f1= λ-> one(λ);
julia> f2= λ-> λ;
julia> f3= λ-> exp(sin(λ/2));
julia> nep=SPMF_NEP([A,B,C],[f1,f2,f3]);
julia> cp=CORKPencil(nep,NleigsCorkLinearization());
julia> (A,B)=buildPencil(cp) # Form the pencil
julia> λv=eigen(A,B).values;
julia> λ=λv[sortperm(abs.(λv))[1]]; # Smallest eigval
julia> minimum(svdvals(compute_Mder(nep,λ))) # It's a solution
2.4364382475487156e-11
```

# References:

*  Compact rational Krylov methods for nonlinear eigenvalue problems SIAM Journal on Matrix Analysis and Applications, 36 (2), 820-838, 2015.

    """
struct CORKPencil{T1<:AbstractMatrix,T2<:AbstractMatrix}
        M::T1
        N::T1
        Av::Vector{T2}   # Array of Array of matrices
        Bv::Vector{T2}   # Array of Array of matrices
    end

    abstract type CorkLinearization end;

    struct DefaultCorkLinearization <: CorkLinearization end;
    # TODO: define default

    struct IarCorkLinearization <: CorkLinearization
        d::Int
        function IarCorkLinearization(;d=10)
            return new(d);
        end
    end

    struct NleigsCorkLinearization <: CorkLinearization
        Σ::Vector
        Ξ::Vector
        maxdgr::Int
        tollin::Float64;
        function NleigsCorkLinearization(;Σ=[-1.0-1im,-1+1im,+1+1im,1-1im],Ξ = [Inf], maxdgr=100 , tollin=1e-6 )
            return new(Σ,Ξ,maxdgr,tollin);
        end
    end

    function CORKPencil(nep,is::IarCorkLinearization)
        M=diagm( 0 =>  ones(is.d) )[2:end,:]
        N=diagm( -1 =>  1 ./ (1:is.d-1) )[2:end,:]
        Av=Array{AbstractMatrix,1}(undef, is.d)
        Bv=Array{AbstractMatrix,1}(undef, is.d)
        Av[1]=-compute_Mder(nep,0,0)
        for j=2:is.d Av[j]=zero(Av[1])              end
        for j=1:is.d Bv[j]=compute_Mder(nep,0,j)/j  end
        return CORKPencil(M,N,Av,Bv)
    end

    function CORKPencil(nep,is::NleigsCorkLinearization)
        D,β,ξ,σ=nleigs_coefficients(nep,is.Σ,tollin=is.tollin,Ξ=is.Ξ)
        d=length(β)-1
        σ=σ[1:d+1]; β=β[1:d+1]; ξ=ξ[1:d+1]
        σ,β,ξ=promote(σ,β,ξ) # be sure that M and N have the same type
        M=diagm( -1 => σ[1:d], 0 =>  β[1:d] )[2:end-1,1:end-1]
        N=diagm( -1 => ones(d), 0 =>  β[1:d]./ξ[1:d] )[2:end-1,1:end-1]
        Av=Array{AbstractMatrix,1}(undef, d)
        Av[1:d-1]=D[1:d-1]; Av[d]=D[d]-σ[d]/β[d+1]*D[d+1]
        Bv=Array{AbstractMatrix,1}(undef, d)
        Bv[1:d-1]=D[1:d-1]/ξ[d+1]; Bv[d]=D[d]/ξ[d+1]-D[d+1]/β[d+1]
        return CORKPencil(M,N,Av,Bv)
    end

    # construct the linearization
"""
    (A,B)=buildPencil(cp)

Constructs a pencil from a `CORKPencil` or `CORKPencilLR`. The returned
matrices correspond to a generalized eigenvalue problem. See also [`CORKPencil`](@ref), [`CORKPencilLR`](@ref). See [`lowRankCompress`] for examples.

"""
    function buildPencil(cp::CORKPencil)
        n=size(cp.Av[1],1)
        if issparse(cp.Av[1])   II=sparse(I, n, n)
        else                    II=Matrix(I, n, n)    end
        return [hcat(cp.Av...); kron(cp.M,II)], [hcat(cp.Bv...); kron(cp.N,II)]
    end

"""
    struct CORKPencilLR
    function CORKPencilLR(M,N,Av,AvLR,Bv,BvLR,Z);

Represents / constructs a low-rank CORK-pencil. `AvLR`, `BvLR`
and `Z` correspond to the low-rank factorization of terms
`Av` and `Bv`. See [`CORKPencil`](@ref) and reference below.

# Example

The example illustrate a low-rank linearization of a Taylor expansion
of the NEP ``A0-λI+vv^Te^{-λ}``.

```julia-repl
julia> A0=[1.0 3.0; -1.0 2.0]/10;
julia> v=reshape([-1.0 ; 1]/sqrt(2),n,1);

julia> Av=[-A0-v*v']
julia> Bv=[-one(A0)-v*v']
julia> BvLR=[v/2, -v/3, v/4, -v/5, v/6, -v/7,  v/8, -v/9]
julia> AvLR=zero.(BvLR);
julia> Z=v;
julia> d=9;
julia> M=diagm( 0 =>  ones(d) )[2:end,:]
julia> N=diagm( -1 =>  1 ./ (1:d-1) )[2:end,:]
julia> cplr=CORKPencilLR(M,N,Av,AvLR,Bv,BvLR,Z);
julia> (AA,BB)=buildPencil(cplr);
julia> λ=eigen(AA,BB).values[end];
julia> minimum(svdvals(A0-λ*I+v*v'*exp(-λ)))
8.870165379112754e-13
```

# References:

*  Section 7 in Compact rational Krylov methods for nonlinear eigenvalue problems SIAM Journal on Matrix Analysis and Applications, 36 (2), 820-838, 2015.
"""
    struct CORKPencilLR
    	M::AbstractMatrix
    	N::AbstractMatrix
    	Av::Vector{AbstractMatrix}
    	AvLR::Vector{AbstractMatrix} # In CORK-paper Avtilde
    	Bv::Vector{AbstractMatrix}
    	BvLR::Vector{AbstractMatrix} # In CORK-paper Bvtilde
        Z::AbstractMatrix
    end

"""
    cplr=lowRankCompress(cp_org::CORKPencil,dtilde,rk)

Constructs a [`CORKPencilLR`](@ref) from a [`CORKPencil`](@ref). This is done by
assuming that terms higher than `dtilde` are of low rank, with rank `rk`.
More precisely, all `A[j]` and `B[j]` for `j>dtilde` are assumed
to be of the form ``C_jZ^T``.

# Example

This illustrates how to form a `CORKPencil` from a NEP, and
subsequently form a smaller pencil using `CORKPencilLR`.

```julia-repl
julia> A0=[1.0 3.0; -1.0 2.0]/10;
julia> v=[-1.0 ; 1]/sqrt(2);
julia> nep=DEP([A0,v*v']);
julia> cp_org=CORKPencil(nep,IarCorkLinearization(d=10));
julia> cplr=lowRankCompress(cp_org::CORKPencil,1,1);
julia> (AA,BB)=buildPencil(cplr);
julia> λ=eigen(AA,BB).values[end];
julia> minimum(svdvals(A0-λ*I+v*v'*exp(-λ))) # Check if it is an eigval
2.7621446071952216e-14
```

"""
function lowRankCompress(cp_org::CORKPencil,dtilde,rk)
    d=length(cp_org.Av);
    # Only use the first low-rank term as a basis (this
    # can cause if the first low-rank term has lower rank
    # than the other ones).
    Z=svd(cp_org.Bv[dtilde+1]).V[:,1:rk];

    # Check that the M and N-matrices have the block triangular structure
    if ((norm(cp_org.M[1:(dtilde-1),(dtilde+1):end]) > 0) ||
        (norm(cp_org.N[1:(dtilde-1),(dtilde+1):end]) > 0))
        error("The M-matrix does not have the required structure. Try increasing dtilde.");
    end

    # Manually extract the low-rank coefficients by
    # multiplying by Z. Works since Z is orthogonal.
    Bvtilde=map(i-> cp_org.Bv[i]*Z,  (dtilde+1):d);
    Avtilde=map(i-> cp_org.Av[i]*Z,  (dtilde+1):d);

    return CORKPencilLR(cp_org.M,cp_org.N,
                        cp_org.Av[1:dtilde],Avtilde,
                        cp_org.Bv[1:dtilde],Bvtilde,
                        Z);
end
function buildPencil(cp::CORKPencilLR)

    # Extract the variables
    n=size(cp.Av[1],1)
    dtilde=length(cp.Av); # Nof full rank terms
    d=length(cp.Av)+length(cp.AvLR);
    rk=size(cp.Z,2);

    # Type logic. The pencil will be of the largest type of the eltypes
    T=promote_type(eltype(cp.Z),
                   eltype(cp.Av[1]),
                   eltype(cp.AvLR[1]),
                   eltype(cp.Bv[1]),
                   eltype(cp.BvLR[1]),
                   eltype(cp.M),
                   eltype(cp.N));

    In=Matrix{T}(I,n,n);
    Idtilde=Matrix{T}(I,rk,rk);

    ## Note: this info is not in CORK paper:
    ##    M11: (dtilde-1) x (dtilde)
    ##    M21: (d-dtilde) x (dtilde)
    ##    M22: (d-dtilde) x (d-dtilde)
    ## same for N

    # Extract Mij and Nij:
    M11=cp.M[1:(dtilde-1),1:dtilde];
    M21=cp.M[(dtilde):end,1:dtilde]
    M22=cp.M[(dtilde):end,(dtilde+1):end]

    N11=cp.N[1:(dtilde-1),1:dtilde];
    N21=cp.N[(dtilde):end,1:dtilde]
    N22=cp.N[(dtilde):end,(dtilde+1):end]

    # Build the B-matrix
    Btilde1=[hcat(cp.Bv[1:dtilde]...) hcat(cp.BvLR...) ]
    Btilde2=[kron(N11,In) zeros((dtilde-1)*n,(d-dtilde)*rk)]
    Btilde3=[kron(N21,cp.Z') kron(N22,Idtilde)];
    Btilde=vcat(Btilde1,Btilde2,Btilde3);

    # Build the A-matrix
    Atilde1=[hcat(cp.Av[1:dtilde]...) hcat(cp.AvLR...) ]
    Atilde2=[kron(M11,In) zeros((dtilde-1)*n,(d-dtilde)*rk)]
    Atilde3=[kron(M21,cp.Z') kron(M22,Idtilde)];
    Atilde=vcat(Atilde1,Atilde2,Atilde3);

    return (Atilde,Btilde);
end

end  # End Module


using Test, LinearAlgebra, Distributions, AdvancedHMC, Random
using AdvancedHMC.Adaptation: WelfordVar, NaiveVar, WelfordCov, NaiveCov, add_sample!, get_var, get_cov, reset!
using DiffResults: GradientResult, value, gradient
using ForwardDiff: gradient!

# Check that the estimated variance is approximately correct.
@testset "Online v.s. naive v.s. true var/cov estimation" begin
    D = 10
    n_samples = 100_000

    var_welford = WelfordVar(D)
    var_naive = NaiveVar()
    var_estimators = [var_welford, var_naive]
    cov_welford = WelfordCov(D)
    cov_naive = NaiveCov()
    cov_estimators = [cov_welford, cov_naive]
    estimators = [var_estimators..., cov_estimators...]

    for dist in [MvNormal(zeros(D), ones(D)), Dirichlet(D, 1)]
        for _ = 1:n_samples
            s = rand(dist)
            for estimator in estimators
                add_sample!(estimator, s)
            end
        end

        @test get_var(var_welford) ≈ get_var(var_naive) atol=0.1D
        for estimator in var_estimators
            @test get_var(estimator) ≈ var(dist) atol=0.1D
        end

        @test get_cov(cov_welford) ≈ get_cov(cov_naive) atol=0.1D^2
        for estimator in cov_estimators
            @test get_cov(estimator) ≈ cov(dist) atol=0.1D^2
        end

        for estimator in estimators
            reset!(estimator)
        end
    end
end

@testset "Preconditioner constructors" begin
    θ = [0.0, 0.0, 0.0, 0.0]
    pc1 = Preconditioner(UnitEuclideanMetric) # default dim = 2
    pc2 = Preconditioner(DiagEuclideanMetric)
    pc3 = Preconditioner(DenseEuclideanMetric)

    # Var adaptor dimention should be increased to length(θ) from 2
    AdvancedHMC.adapt!(pc1, θ, 1.)
    AdvancedHMC.adapt!(pc2, θ, 1.)
    AdvancedHMC.adapt!(pc3, θ, 1.)
    @test AdvancedHMC.Adaptation.getM⁻¹(pc2) == ones(length(θ))
    @test AdvancedHMC.Adaptation.getM⁻¹(pc3) == LinearAlgebra.diagm(0 => ones(length(θ)))
end

@testset "Stan HMC adaptors" begin
    θ = [0.0, 0.0, 0.0, 0.0]

    adaptor1 = StanHMCAdaptor(
        1000,
        Preconditioner(UnitEuclideanMetric),
        NesterovDualAveraging(0.8, 0.5)
    )
    adaptor2 = StanHMCAdaptor(
        1000,
        Preconditioner(DiagEuclideanMetric),
        NesterovDualAveraging(0.8, 0.5)
    )
    adaptor3 = StanHMCAdaptor(
        1000,
        Preconditioner(DenseEuclideanMetric),
        NesterovDualAveraging(0.8, 0.5)
    )

    AdvancedHMC.adapt!(adaptor1, θ, 1.)
    AdvancedHMC.adapt!(adaptor2, θ, 1.)
    AdvancedHMC.adapt!(adaptor3, θ, 1.)
    @test AdvancedHMC.Adaptation.getM⁻¹(adaptor2) == ones(length(θ))
    @test AdvancedHMC.Adaptation.getM⁻¹(adaptor3) == LinearAlgebra.diagm(0 => ones(length(θ)))
end

let D=10
    function buildfuncs(target)
        ℓπ(θ) = logpdf(target, θ)

        function ∂ℓπ∂θ(θ)
            res = GradientResult(θ)
            gradient!(res, ℓπ, θ)
            return (value(res), gradient(res))
        end

        return ℓπ, ∂ℓπ∂θ
    end

    function runnuts(ℓπ, ∂ℓπ∂θ, metric; n_samples=2_000)
        n_adapts = 1_000

        θ_init = randn(D)

        h = Hamiltonian(metric, ℓπ, ∂ℓπ∂θ)
        prop = NUTS(Leapfrog(find_good_eps(h, θ_init)))
        adaptor = StanHMCAdaptor(
            n_adapts,
            Preconditioner(metric),
            NesterovDualAveraging(0.8, prop.integrator)
        )
        samples, stats = sample(h, prop, θ_init, n_samples, adaptor, n_adapts; verbose=false)
        return (samples=samples, stats=stats, adaptor=adaptor)
    end

    @testset "Adapted mass v.s. true variance" begin
        Random.seed!(123)
        n_tests = 5

        @testset "DiagEuclideanMetric" begin
            for _ in 1:n_tests
                # Random variance
                σ = ones(D) + abs.(randn(D))

                # Diagonal Gaussian
                target = MvNormal(zeros(D), σ)
                ℓπ, ∂ℓπ∂θ = buildfuncs(target)

                res = runnuts(ℓπ, ∂ℓπ∂θ, DiagEuclideanMetric(D))
                @test res.adaptor.pc.var ≈ σ .^ 2 rtol=0.2

                res = runnuts(ℓπ, ∂ℓπ∂θ, DenseEuclideanMetric(D))
                @test res.adaptor.pc.covar ≈ diagm(0 => σ.^2) rtol=0.25
            end
        end

        @testset "DenseEuclideanMetric" begin
            for _ in 1:n_tests
                # Random covariance
                m = randn(D, D)
                Σ = m' * m

                # Correlated Gaussian
                target = MvNormal(zeros(D), Σ)
                ℓπ, ∂ℓπ∂θ = buildfuncs(target)

                res = runnuts(ℓπ, ∂ℓπ∂θ, DiagEuclideanMetric(D))
                @test res.adaptor.pc.var ≈ diag(Σ) rtol=0.2

                res = runnuts(ℓπ, ∂ℓπ∂θ, DenseEuclideanMetric(D))
                @test res.adaptor.pc.covar ≈ Σ rtol=0.25
            end
        end

    end

    @testset "Initialisation adaptor by metric" begin
        target = MvNormal(zeros(D), 1)
        ℓπ, ∂ℓπ∂θ = buildfuncs(target)

        mass_init = fill(0.5, D)
        res = runnuts(ℓπ, ∂ℓπ∂θ, DiagEuclideanMetric(mass_init); n_samples=1)
        @test res.adaptor.pc.var == mass_init

        mass_init = diagm(0 => fill(0.5, D))
        res = runnuts(ℓπ, ∂ℓπ∂θ, DenseEuclideanMetric(mass_init); n_samples=1)
        @test res.adaptor.pc.covar == mass_init
    end
end


# Test sample functions

using StatsBase
using Base.Test

srand(1234)

function est_p(x, K)
	h = zeros(Int, K)
	for xi in x
		h[xi] += 1
	end
	p = h / length(x)
end

#### randi

const randi = StatsBase.randi
n = 1_000_000

x = [randi(10) for i = 1:n]
@test isa(x, Vector{Int})
@test minimum(x) == 1
@test maximum(x) == 10

x = [randi(3, 12) for i = 1:n]
@test isa(x, Vector{Int})
@test minimum(x) == 3
@test maximum(x) == 12

#### sample with replacement

## individual

x = [sample(1:5) for _ = 1:10^5]
p = fill(0.2, 5)
@test isa(x, Vector{Int})
@test_approx_eq_eps est_p(x, 5) p 0.02

## unordered

x = sample(1:5, 10^5)
@test isa(x, Vector{Int})
@test length(x) == 10^5
@test !issorted(x)  # extremely unlikely to be sorted for n=10^5
@test_approx_eq_eps est_p(x, 5) p 0.02
@test_approx_eq_eps est_p(x[1:50000], 5) p 0.03

## ordered

x = sample(1:5, 10^5; ordered=true)
@test isa(x, Vector{Int})
@test length(x) == 10^5
@test issorted(x)
@test_approx_eq_eps est_p(x, 5) p 0.02

#### sample without replacement

## unordered (using samplepair)

x = Array(Int, 2, n)
for i = 1:n
	x[:,i] = sample(1:5, 2; replace=false)
end

@test all(x[1,:] .!= x[2,:])
@test_approx_eq_eps est_p(x[1,:], 5) p 0.02
@test_approx_eq_eps est_p(x[2,:], 5) p 0.02

## unordered (using Fisher-Yates)

x = Array(Int, 3, n)
for i = 1:n
	x[:,i] = sample(1:5, 3; replace=false)
end

@test all(x[1,:] .!= x[2,:])
@test all(x[1,:] .!= x[3,:])
@test all(x[2,:] .!= x[3,:])
@test_approx_eq_eps est_p(x[1,:], 5) p 0.02
@test_approx_eq_eps est_p(x[2,:], 5) p 0.02
@test_approx_eq_eps est_p(x[3,:], 5) p 0.02

## unordered (using Self-avoid method)

x = Array(Int, 3, 10000)
for i = 1:10000
	x[:,i] = sample(1:1000, 3; replace=false)
end

@test all(x[1,:] .!= x[2,:])
@test all(x[1,:] .!= x[3,:])
@test all(x[2,:] .!= x[3,:])

## ordered

x = Array(Int, 3, n)
for i = 1:n
	x[:,i] = sample(1:5, 3; replace=false, ordered=true)
end

@test all(x[1,:] .< x[2,:])
@test all(x[2,:] .< x[3,:])
@test_approx_eq_eps est_p(vec(x), 5) p 0.02

#### weighted sample with replacement

wv = weights([1.0, 2.0, 4.5, 2.5])
p = [0.10, 0.20, 0.45, 0.25]

## individual

x = Int[sample(wv) for _ = 1:10^5]
@test isa(x, Vector{Int})
@test_approx_eq_eps est_p(x, 4) p 0.02

## unordered

x = sample(1:5, wv, 10^5)
@test isa(x, Vector{Int})
@test length(x) == 10^5
@test !issorted(x)
@test_approx_eq_eps est_p(x, 4) p 0.02

## ordered

x = sample(1:5, wv, 10^5; ordered=true)
@test isa(x, Vector{Int})
@test length(x) == 10^5
@test issorted(x)
@test_approx_eq_eps est_p(x, 4) p 0.02


using Mamba

## Data
rats = Dict{Symbol, Any}(
  :y =>
    [151, 199, 246, 283, 320,
     145, 199, 249, 293, 354,
     147, 214, 263, 312, 328,
     155, 200, 237, 272, 297,
     135, 188, 230, 280, 323,
     159, 210, 252, 298, 331,
     141, 189, 231, 275, 305,
     159, 201, 248, 297, 338,
     177, 236, 285, 350, 376,
     134, 182, 220, 260, 296,
     160, 208, 261, 313, 352,
     143, 188, 220, 273, 314,
     154, 200, 244, 289, 325,
     171, 221, 270, 326, 358,
     163, 216, 242, 281, 312,
     160, 207, 248, 288, 324,
     142, 187, 234, 280, 316,
     156, 203, 243, 283, 317,
     157, 212, 259, 307, 336,
     152, 203, 246, 286, 321,
     154, 205, 253, 298, 334,
     139, 190, 225, 267, 302,
     146, 191, 229, 272, 302,
     157, 211, 250, 285, 323,
     132, 185, 237, 286, 331,
     160, 207, 257, 303, 345,
     169, 216, 261, 295, 333,
     157, 205, 248, 289, 316,
     137, 180, 219, 258, 291,
     153, 200, 244, 286, 324],
  :x => [8.0, 15.0, 22.0, 29.0, 36.0]
)
rats[:xbar] = mean(rats[:x])
rats[:N] = size(rats[:y], 1)
rats[:T] = size(rats[:y], 2)

rats[:rat] = Int[div(i - 1, 5) + 1 for i in 1:150]
rats[:week] = Int[(i - 1) % 5 + 1 for i in 1:150]
rats[:X] = rats[:x][rats[:week]]
rats[:Xm] = rats[:X] - rats[:xbar]

## Model Specification
model = Model(

  y = Stochastic(1,
    (alpha, beta, rat, Xm, s2_c) ->
      begin
        mu = alpha[rat] + beta[rat] .* Xm
        MvNormal(mu, sqrt(s2_c))
      end,
    false
  ),

  alpha = Stochastic(1,
    (mu_alpha, s2_alpha) -> Normal(mu_alpha, sqrt(s2_alpha)),
    false
  ),

  alpha0 = Logical(
    (mu_alpha, xbar, mu_beta) -> mu_alpha - xbar * mu_beta
  ),

  mu_alpha = Stochastic(
    () -> Normal(0.0, 1000),
    false
  ),

  s2_alpha = Stochastic(
    () -> InverseGamma(0.001, 0.001),
    false
  ),

  beta = Stochastic(1,
    (mu_beta, s2_beta) -> Normal(mu_beta, sqrt(s2_beta)),
    false
  ),

  mu_beta = Stochastic(
    () -> Normal(0.0, 1000)
  ),

  s2_beta = Stochastic(
    () -> InverseGamma(0.001, 0.001),
    false
  ),

  s2_c = Stochastic(
    () -> InverseGamma(0.001, 0.001)
  )

)

## Initial Values
inits = [
  Dict(:y => rats[:y], :alpha => fill(250, 30), :beta => fill(6, 30),
       :mu_alpha => 150, :mu_beta => 10, :s2_c => 1, :s2_alpha => 1,
       :s2_beta => 1),
  Dict(:y => rats[:y], :alpha => fill(20, 30), :beta => fill(0.6, 30),
       :mu_alpha => 15, :mu_beta => 1, :s2_c => 10, :s2_alpha => 10,
       :s2_beta => 10)
]

## Sampling Scheme
scheme = [Slice(:s2_c, 10.0),
          AMWG(:alpha, 100.0),
          Slice([:mu_alpha, :s2_alpha], [100.0, 10.0], Univariate),
          AMWG(:beta, 1.0),
          Slice([:mu_beta, :s2_beta], 1.0, Univariate)]
setsamplers!(model, scheme)

## MCMC Simulations
sim = mcmc(model, rats, inits, 10000, burnin=2500, thin=2, chains=2)
describe(sim)


const TEST_FLOAT32 = false
const TEST_DOUBLE64 = false
const TEST_STATIC_SIZED = false

using Manifolds
using ManifoldsBase

using LinearAlgebra
using Distributions
using DoubleFloats
using ForwardDiff
using Random
using ReverseDiff
using StaticArrays
using Test
using LightGraphs
using SimpleWeightedGraphs

find_eps(x::Type{TN}) where TN<:Number = eps(real(TN))
find_eps(x) = find_eps(number_eltype(x))
find_eps(x...) = find_eps(Base.promote_type(map(number_eltype, x)...))

"""
    test_manifold(m::Manifold, pts::AbstractVector;
    args
    )

Tests general properties of manifold `m`, given at least three different points
that lie on it (contained in `pts`).

# Arguments
- `default_inverse_retraction_method = ManifoldsBase.LogarithmicInverseRetraction()` - default method for inverse retractions ([`log`](@ref))
- `default_retraction_method = ManifoldsBase.ExponentialRetraction()` - default method for retractions ([`exp`](@ref))
- `exp_log_atol_multiplier = 0`, change absolute tolerance of exp/log tests (0 use default, i.e. deactivate atol and use rtol)
- `exp_log_rtol_multiplier = 1`, change the relative tolerance of exp/log tests (1 use default). This is deactivated if the `exp_log_atol_multiplier` is nonzero.
- `retraction_atol_multiplier = 0`, change absolute tolerance of (inverse) retraction tests (0 use default, i.e. deactivate atol and use rtol)
- `retraction_rtol_multiplier = 1`, change the relative tolerance of (inverse) retraction tests (1 use default). This is deactivated if the `exp_log_atol_multiplier` is nonzero.
- `inverse_retraction_methods = []`: inverse retraction methods that will be tested.
- `point_distributions = []` : point distributions to test
- `projection_tvector_atol_multiplier = 0` : chage absolute tolerance in testing projections (0 use default, i.e. deactivate atol and use rtol)
-  tvector_distributions = []` : tangent vector distributions to test
- `basis_types = ()` : basis types that will be tested
- `rand_tvector_atol_multiplier = 0` : chage absolute tolerance in testing random vectors (0 use default, i.e. deactivate atol and use rtol)
  random tangent vectors are tangent vectors
- `retraction_methods = []`: retraction methods that will be tested.
- `test_forward_diff = true`: if true, automatic differentiation using
  ForwardDiff is tested.
- `test_reverse_diff = true`: if true, automatic differentiation using
  ReverseDiff is tested.
- `test_musical_isomorphisms = false` : test musical isomorphisms
- `test_mutating_rand = false` : test the mutating random function for points on manifolds
- `test_project_tangent = false` : test projections on tangent spaces
- `test_representation_size = true` : test repersentation size of points/tvectprs
- `test_tangent_vector_broadcasting = true` : test boradcasting operators on TangentSpace
- `test_vector_transport = false` : test vector transport
- `test_vector_spaces = true` : test Vector bundle of this manifold
"""
function test_manifold(M::Manifold, pts::AbstractVector;
    test_exp_log = true,
    test_is_tangent = true,
    test_injectivity_radius=true,
    test_forward_diff = true,
    test_reverse_diff = true,
    test_tangent_vector_broadcasting = true,
    test_project_tangent = false,
    test_project_point = false,
    test_representation_size = true,
    test_musical_isomorphisms = false,
    test_vector_transport = false,
    test_mutating_rand = false,
    test_vector_spaces = true,
    test_vee_hat = false,
    is_mutating = true,
    default_inverse_retraction_method = ManifoldsBase.LogarithmicInverseRetraction(),
    default_retraction_method = ManifoldsBase.ExponentialRetraction(),
    retraction_methods = [],
    inverse_retraction_methods = [],
    point_distributions = [],
    tvector_distributions = [],
    basis_types_vecs = (),
    basis_types_to_from = (),
    basis_has_specialized_diagonalizing_get = false,
    exp_log_atol_multiplier = 0,
    exp_log_rtol_multiplier = 1,
    retraction_atol_multiplier = 0,
    retraction_rtol_multiplier = 1,
    projection_atol_multiplier = 0,
    rand_tvector_atol_multiplier = 0,
    is_tangent_atol_multiplier=1,
)

    length(pts) ≥ 3 || error("Not enough points (at least three expected)")
    isapprox(M, pts[1], pts[2]) && error("Points 1 and 2 are equal")
    isapprox(M, pts[1], pts[3]) && error("Points 1 and 3 are equal")

    # get a default tangent vector for every of the three tangent spaces
    n = length(pts)
    if default_inverse_retraction_method === nothing
        tv = [ zero_tangent_vector(M,pts[i]) for i=1:n ] # no other available
    else
        tv = [ inverse_retract(M,pts[i],pts[((i+1)%n)+1],default_inverse_retraction_method) for i=1:n]
    end
    @testset "dimension" begin
        @test isa(manifold_dimension(M), Integer)
        @test manifold_dimension(M) ≥ 0
        @test manifold_dimension(M) == vector_space_dimension(Manifolds.VectorBundleFibers(Manifolds.TangentSpace, M))
        @test manifold_dimension(M) == vector_space_dimension(Manifolds.VectorBundleFibers(Manifolds.CotangentSpace, M))
    end

    test_representation_size && @testset "representation" begin
        function test_repr(repr)
            @test isa(repr, Tuple)
            for rs ∈ repr
                @test rs > 0
            end
        end

        test_repr(Manifolds.representation_size(M))
        for fiber ∈ (Manifolds.TangentSpace, Manifolds.CotangentSpace)
            test_repr(Manifolds.representation_size(Manifolds.VectorBundleFibers(fiber, M)))
        end
    end

    test_injectivity_radius && @testset "injectivity radius" begin
        @test injectivity_radius(M, pts[1]) > 0
        @test injectivity_radius(M, pts[1]) ≥ injectivity_radius(M)
        for rm ∈ retraction_methods
            @test injectivity_radius(M, rm) > 0
            @test injectivity_radius(M, pts[1], rm) ≥ injectivity_radius(M, rm)
            @test injectivity_radius(M, pts[1], rm) ≤ injectivity_radius(M, pts[1])
        end
    end

    @testset "is_manifold_point" begin
        for pt ∈ pts
            @test is_manifold_point(M, pt)
            @test check_manifold_point(M, pt) === nothing
            @test check_manifold_point(M, pt) === nothing
        end
    end

    test_is_tangent && @testset "is_tangent_vector" begin
        for (p, X) in zip(pts, tv)
            atol = is_tangent_atol_multiplier * find_eps(p)
            if !( check_tangent_vector(M, p, X; atol = atol) === nothing )
                print( check_tangent_vector(M, p, X; atol = atol) )
            end
            @test is_tangent_vector(M, p, X; atol = atol)
            @test check_tangent_vector(M, p, X; atol = atol) === nothing
        end
    end

    test_exp_log && @testset "log/exp tests" begin
        epsp1p2 = find_eps(pts[1], pts[2])
        atolp1p2 = exp_log_atol_multiplier * epsp1p2
        rtolp1p2 = exp_log_atol_multiplier == 0. ? sqrt(epsp1p2)*exp_log_rtol_multiplier : 0
        X1 = log(M, pts[1], pts[2])
        X2 = log(M, pts[2], pts[1])
        @test isapprox(M, pts[2], exp(M, pts[1], X1); atol = atolp1p2, rtol = rtolp1p2)
        @test isapprox(M, pts[1], exp(M, pts[1], X1, 0); atol = atolp1p2, rtol = rtolp1p2)
        @test isapprox(M, pts[2], exp(M, pts[1], X1, 1); atol = atolp1p2, rtol = rtolp1p2)
        @test isapprox(M, pts[1], exp(M, pts[2], X2); atol = atolp1p2, rtol = rtolp1p2)
        @test is_manifold_point(M, exp(M, pts[1], X1); atol = atolp1p2, rtol = rtolp1p2)
        @test isapprox(M, pts[1], exp(M, pts[1], X1, 0); atol = atolp1p2, rtol = rtolp1p2)
        for p ∈ pts
            epsx = find_eps(p)
            @test isapprox(M, p, zero_tangent_vector(M, p), log(M, p, p);
                atol = epsx * exp_log_atol_multiplier,
                rtol = exp_log_atol_multiplier == 0. ? sqrt(epsx)*exp_log_rtol_multiplier : 0
            )
            @test isapprox(M, p, zero_tangent_vector(M, p), inverse_retract(M, p, p);
                atol = epsx * exp_log_atol_multiplier,
                rtol = exp_log_atol_multiplier == 0. ? sqrt(epsx)*exp_log_rtol_multiplier : 0.
            )
        end
        atolp1 = exp_log_atol_multiplier * find_eps(pts[1])
        if is_mutating
            zero_tangent_vector!(M, X1, pts[1])
        else
            X1 = zero_tangent_vector(M, pts[1])
        end
        @test isapprox(M, pts[1], X1, zero_tangent_vector(M, pts[1]); atol = atolp1)
        if is_mutating
            log!(M, X1, pts[1], pts[2])
        else
            X1 = log(M,pts[1], pts[2])
        end

        @test isapprox(M, exp(M, pts[1], X1, 1), pts[2]; atol = atolp1)
        @test isapprox(M, exp(M, pts[1], X1, 0), pts[1]; atol = atolp1)

        @test distance(M, pts[1], pts[2]) ≈ norm(M, pts[1], X1)

        X3 = log(M, pts[1], pts[3])

        @test real(inner(M, pts[1], X1, X3)) ≈ real(inner(M, pts[1], X3, X1))
        @test imag(inner(M, pts[1], X1, X3)) ≈ -imag(inner(M, pts[1], X3, X1))
        @test imag(inner(M, pts[1], X1, X1)) ≈ 0

        @test norm(M, pts[1], X1) isa Real
        @test norm(M, pts[1], X1) ≈ sqrt(inner(M, pts[1], X1, X1))
    end

    @testset "(inverse &) retraction tests" begin
        for (p, X) in zip(pts, tv)
            epsx = find_eps(p)
            for retr_method ∈ retraction_methods
                @test is_manifold_point(M, retract(M, p, X, retr_method))
                @test isapprox(M, p, retract(M, p, X, 0, retr_method);
                    atol = epsx * retraction_atol_multiplier,
                    rtol = retraction_atol_multiplier == 0 ? sqrt(epsx)*retraction_rtol_multiplier : 0
                )
                if is_mutating
                    new_pt = allocate(p)
                    retract!(M, new_pt, p, X, retr_method)
                else
                    new_pt = retract(M, p, X, retr_method)
                end
                @test is_manifold_point(M, new_pt)
            end
        end
        for p ∈ pts
            epsx = find_eps(p)
            for inv_retr_method ∈ inverse_retraction_methods
                @test isapprox(M, p, zero_tangent_vector(M, p), inverse_retract(M, p, p, inv_retr_method);
                    atol = epsx * retraction_atol_multiplier,
                    rtol = retraction_atol_multiplier == 0 ? sqrt(epsx)*retraction_rtol_multiplier : 0
                )
            end
        end
    end

    test_vector_spaces && @testset "vector spaces tests" begin
        for p ∈ pts
            X = zero_tangent_vector(M, p)
            mts = Manifolds.VectorBundleFibers(Manifolds.TangentSpace, M)
            @test isapprox(M, p, X, zero_vector(mts, p))
            zero_vector!(mts, X, p)
            @test isapprox(M, p, X, zero_tangent_vector(M,p))
        end
    end

    @testset "basic linear algebra in tangent space" begin
        for (p, X) in zip(pts, tv)
            @test isapprox(M, p, 0*X, zero_tangent_vector(M, p); atol = find_eps(pts[1]))
            @test isapprox(M, p, 2*X, X+X)
            @test isapprox(M, p, 0*X, X-X)
            @test isapprox(M, p, (-1)*X, -X)
        end
    end

    test_tangent_vector_broadcasting && @testset "broadcasted linear algebra in tangent space" begin
        for (p, X) in zip(pts, tv)
            @test isapprox(M, p, 3*X, 2 .* X .+ X)
            @test isapprox(M, p, -X, X .- 2 .* X)
            @test isapprox(M, p, -X, .-X)
            if (isa(X, AbstractArray))
                Y = allocate(X)
                Y .= 2 .* X .+ X
            else
                Y = 2*X+X
            end
            @test isapprox(M, p, Y, 3*X)
        end
    end

    test_project_tangent && @testset "project tangent test" begin
        for (p, X) in zip(pts, tv)
            atol = find_eps(p) * projection_atol_multiplier
            @test isapprox(M, p, X, project(M, p, X); atol = atol)
            if is_mutating
                X2 = allocate(X)
                project!(M, X2, p, X)
            else
                X2 = project(M, p, X)
            end
            @test isapprox(M, p, X2, X; atol = atol)
        end
    end

    test_project_point && @testset "project point test" begin
        for p in pts
            atol = find_eps(p) * projection_atol_multiplier
            @test isapprox(M, p, project(M, p); atol = atol)
            if is_mutating
                p2 = allocate(p)
                project!(M, p2, p)
            else
                p2 = project(M, p)
            end
            @test isapprox(M, p2, p; atol = atol)
        end
    end

    test_vector_transport && !( default_inverse_retraction_method === nothing) && @testset "vector transport" begin
        tvatol = is_tangent_atol_multiplier*find_eps(pts[1])
        X1 = inverse_retract(M, pts[1], pts[2], default_inverse_retraction_method)
        X2 = inverse_retract(M, pts[1], pts[3], default_inverse_retraction_method)
        v1t1 = vector_transport_to(M, pts[1], X1, pts[3])
        v1t2 = vector_transport_direction(M, pts[1], X1, X2)
        @test is_tangent_vector(M, pts[3], v1t1; atol=tvatol)
        @test is_tangent_vector(M, pts[3], v1t2; atol=tvatol)
        @test isapprox(M, pts[3], v1t1, v1t2)
        @test isapprox(M, pts[1], vector_transport_to(M, pts[1], X1, pts[1]), X1)
    end

    for btype ∈ basis_types_vecs
        @testset "Basis support for $(btype)" begin
            p = pts[1]
            b = get_basis(M, p, btype)
            @test isa(b, CachedBasis)
            bvectors = get_vectors(M, p, b)
            N = length(bvectors)
            @test real_dimension(number_system(btype)) * N == manifold_dimension(M)

            # test orthonormality
            for i in 1:N
                @test norm(M, p, bvectors[i]) ≈ 1
                for j in i+1:N
                    @test real(inner(M, p, bvectors[i], bvectors[j])) ≈ 0 atol = sqrt(find_eps(p))
                end
            end
            if isa(btype, ProjectedOrthonormalBasis)
                # check projection idempotency
                for i in 1:N
                    @test norm(M, p, bvectors[i]) ≈ 1
                    for j in i+1:N
                        @test real(inner(M, p, bvectors[i], bvectors[j])) ≈ 0 atol = sqrt(find_eps(p))
                    end
                end
                # check projection idempotency
                for i in 1:N
                    @test isapprox(M, p, project(M, p, bvectors[i]), bvectors[i])
                end
            end
            if !isa(btype, ProjectedOrthonormalBasis) &&
                (basis_has_specialized_diagonalizing_get || !isa(btype, DiagonalizingOrthonormalBasis))

                X1 = inverse_retract(M, p, pts[2], default_inverse_retraction_method)
                Xb = get_coordinates(M, p, X1, btype)

                @test get_coordinates(M, p, X1, b) ≈ Xb
                @test isapprox(M, p, get_vector(M, p, Xb, b), get_vector(M, p, Xb, btype))
            end
        end
    end

    for btype ∈ (basis_types_to_from..., basis_types_vecs...)
        p = pts[1]
        N = manifold_dimension(M)
        if !isa(btype, ProjectedOrthonormalBasis) &&
            (basis_has_specialized_diagonalizing_get || !isa(btype, DiagonalizingOrthonormalBasis))

            X1 = inverse_retract(M, p, pts[2], default_inverse_retraction_method)

            Xb = get_coordinates(M, p, X1, btype)
            #@test isa(Xb, AbstractVector{<:Real})
            @test length(Xb) == N
            Xbi = get_vector(M, p, Xb, btype)
            @test isapprox(M, p, X1, Xbi)

            Xs = [[ifelse(i==j, 1, 0) for j in 1:N] for i in 1:N]
            Xs_invs = [get_vector(M, p, Xu, btype) for Xu in Xs]
            # check orthonormality of inverse representation
            for i in 1:N
                @test norm(M, p, Xs_invs[i]) ≈ 1 atol = find_eps(p)
                for j in i+1:N
                    @test real(inner(M, p, Xs_invs[i], Xs_invs[j])) ≈ 0 atol = sqrt(find_eps(p))
                end
            end

            if is_mutating
                Xb_s = allocate(Xb)
                @test get_coordinates!(M, Xb_s, p, X1, btype) === Xb_s
                @test isapprox(Xb_s, Xb; atol = find_eps(p))

                Xbi_s = allocate(Xbi)
                @test get_vector!(M, Xbi_s, p, Xb, btype) === Xbi_s
                @test isapprox(M, p, X1, Xbi_s)
            end
        end
    end

    test_vee_hat && @testset "vee and hat" begin
        p = pts[1]
        q = pts[2]
        X = inverse_retract(M, p, q, default_inverse_retraction_method)
        Y = vee(M, p, X)
        @test length(Y) == manifold_dimension(M)
        @test isapprox(M, p, X, hat(M, p, Y))
        Y2 = allocate(Y)
        vee_ret = vee!(M, Y2, p, X)
        @test vee_ret === Y2
        @test isapprox(Y, Y2)
        X2 = allocate(X)
        hat_ret = hat!(M, X2, p, Y)
        @test hat_ret === X2
        @test isapprox(M, p, X2, X)
    end

    test_forward_diff && @testset "ForwardDiff support" begin
        for (p, X) in zip(pts, tv)
            exp_f(t) = distance(M, p, exp(M, p, t[1]*X))
            d12 = norm(M, p, X)
            for t ∈ 0.1:0.1:0.9
                @test d12 ≈ ForwardDiff.derivative(exp_f, t)
            end

            retract_f(t) = distance(M, p, retract(M, p, t[1]*X))
            for t ∈ 0.1:0.1:0.9
                @test ForwardDiff.derivative(retract_f, t) ≥ 0
            end
        end
    end

    test_reverse_diff && @testset "ReverseDiff support" begin
        for (p, X) in zip(pts, tv)
            exp_f(t) = distance(M, p, exp(M, p, t[1]*X))
            d12 = norm(M, p, X)
            for t ∈ 0.1:0.1:0.9
                @test d12 ≈ ReverseDiff.gradient(exp_f, [t])[1]
            end

            retract_f(t) = distance(M, p, retract(M, p, t[1]*X))
            for t ∈ 0.1:0.1:0.9
                @test ReverseDiff.gradient(retract_f, [t])[1] ≥ 0
            end
        end
    end

    test_musical_isomorphisms && @testset "Musical isomorphisms" begin
        if default_inverse_retraction_method !== nothing
            tv_m = inverse_retract(M, pts[1], pts[2],default_inverse_retraction_method)
        else
            tv_m = zero_tangent_vector(M,pts[1])
        end
        ctv_m = flat(M, pts[1], FVector(TangentSpace, tv_m))
        @test ctv_m.type == CotangentSpace
        tv_m_back = sharp(M, pts[1], ctv_m)
        @test tv_m_back.type == TangentSpace
    end

    @testset "number_eltype" begin
        for (p, X) in zip(pts, tv)
            @test number_eltype(X) == number_eltype(p)
            p = retract(M, p, X, default_retraction_method)
            @test number_eltype(p) == number_eltype(p)
        end
    end

    is_mutating && @testset "copyto!" begin
        for (p, X) in zip(pts, tv)
            p2 = allocate(p)
            copyto!(p2, p)
            @test isapprox(M, p2, p)

            X2 = allocate(X)
            if default_inverse_retraction_method === nothing
                X3 = zero_tangent_vector(M, p)
                copyto!(X2, X3)
                @test isapprox(M, p, X2, zero_tangent_vector(M, p))
            else
                q = retract(M, p, X, default_retraction_method)
                X3 = inverse_retract(M, p, q, default_inverse_retraction_method)
                copyto!(X2, X3)
                @test isapprox(M, p, X2, inverse_retract(M, p, q, default_inverse_retraction_method))
            end
        end
    end

    is_mutating && @testset "point distributions" begin
        for p in pts
            prand = allocate(p)
            for pd ∈ point_distributions
                for _ in 1:10
                    @test is_manifold_point(M, rand(pd))
                    if test_mutating_rand
                        rand!(pd, prand)
                        @test is_manifold_point(M, prand)
                    end
                end
            end
        end
    end

    @testset "tangent vector distributions" begin
        for tvd ∈ tvector_distributions
            supp = Manifolds.support(tvd)
            for _ in 1:10
                randtv = rand(tvd)
                atol = rand_tvector_atol_multiplier * find_eps(randtv)
                @test is_tangent_vector(M, supp.point, randtv; atol = atol)
            end
        end
    end
end


@testset "gather" begin
    @testset "3-argumented" begin
        X = [1 2 3; 4 5 6]
        ind1 = [1 2 2; 1 1 1]
        ind2 = [1 2; 3 1]
        for T = [Int8, Int16, Int32, Int64, Int128,
                 Float16, Float32, Float64]
            @test gather(T.(X), ind1, 1) == T.([1 5 6; 1 2 3])
            @test gather(T.(X), ind2, 2) == T.([1 2; 6 4])
        end
    end

    @testset "2-argumented" begin
        input = [3 3 4 4 5;
                 5 5 6 6 7]
        index = [1 2 3 4;
                 4 2 1 3;
                 3 5 5 3]
        output = cat([3 4 4; 5 6 6], [3 3 5; 5 5 7],
                     [4 3 5; 6 5 7], [4 4 4; 6 6 6], dims=3)
        @test gather(input, index) == output
    end
end


type interval

end

function bisection(neg_f_position::Float64, pos_f_position::Float64, func::Function, eps1::Float64)
    return bisection(neg_f_position, pos_f_position, func, eps1, 0.0)
end

function bisection(neg_f_position::Float64, pos_f_position::Float64, func::Function, eps1::Float64, eps2::Float64)
    func_neg_f_position = func(neg_f_position)
    if (func_neg_f_position >= 0.0)
        @show func_neg_f_position
        error("func(neg_f_position) >= 0.0")
    end

    func_pos_f_position = func(pos_f_position)
    if (func_pos_f_position <= 0.0)
        @show func_pos_f_position
        @show func(pos_f_position)
        error("func(pos_f_position) <= 0.0")
    end

    return recursive_bisection(neg_f_position, pos_f_position, func, eps1, eps2)
end

function recursive_bisection(neg_f_position::Float64, pos_f_position::Float64, func::Function, eps1::Float64, eps2::Float64)
    # find zero of increasing func
    # do with for loop?
    centre_point = (pos_f_position + neg_f_position)/2.0
    val = func(centre_point)

    terminate = abs(val) > eps1 && abs(neg_f_position - pos_f_position) > eps2

    if val > 0.0
        pos_f_position = centre_point
    elseif val < 0.0
        neg_f_position = centre_point
    elseif val == 0.0
        return centre_point, centre_point
    end

    if abs(val) < eps1
        return centre_point, centre_point
    elseif abs(neg_f_position - pos_f_position) < eps2
        return neg_f_position, pos_f_position
    else
        return recursive_bisection(neg_f_position, pos_f_position, func, eps1, eps2)
    end
end

if false
    function my_test_func(x::Float64)
        return x
    end

    bisection(-1.0,2.0,my_test_func,0.01)
end


# The Computer Language Benchmarks Game
# https://salsa.debian.org/benchmarksgame-team/benchmarksgame/
#
# Contributed by Adam Beckmeyer

module NBodyAB

import Printf: @printf
import Base: @propagate_inbounds

const DAYS_PER_YEAR = 365.24
const SOLAR_MASS = 4 * π * π
const DEFAULT_STEP = 0.01

# Use a tuple rather than struct to hint compiler towards simd code
# 4 floats in tuple instead of 3 generates better SIMD instructions
V3d(x=0.0, y=0.0, z=0.0) = (Float64(x), Float64(y), Float64(z), 0.0)

# Rust implementation #7 uses pos and v of length 3 with 1 filler
# Float64 field for a total struct width of 64 bytes instead of 72
# bytes as shown currently. This might be possible for Julia in the
# future using https://github.com/eschnett/SIMD.jl. Not sure if it
# would actually be advantageous though.
struct Body
    pos::NTuple{4,Float64}
    v::NTuple{4,Float64}
    m::Float64
end#struct

# don't bother adding the empty 4th element
Base.sum(v::NTuple{4,Float64}) = @inbounds +(v[1], v[2], v[3])

function init_sun(bodies)
    p = V3d()
    for b in bodies
        p = muladd.(b.v, b.m, p)
    end#for
    Body(V3d(), p .* (-inv(SOLAR_MASS)), SOLAR_MASS)
end#function

@propagate_inbounds function update_velocity(b1::Body, b2::Body, Δt::Float64)
    Δpos = b1.pos .- b2.pos
    d² = sum(Δpos .* Δpos)

    # Fastest implementations use an intrinsic to do a
    # single-precision sqrt approximation followed by two iterations
    # of the Newton-Raphson method. But the sqrt isn't currently this
    # implementations bottleneck, so I've left it unoptimized.

    # The fastest Rust/C implementations store this value instead of
    # immediately modifying the body objects. Possible performance
    # gains there due to cache-locality? I've been unable to replicate
    # this in Julia by calculating mag as an NTuple{10,Float64}.
    vmag = Δpos .* (Δt / (d² * √d²))

    (Body(b1.pos, muladd.(-b2.m, vmag, b1.v), b1.m),
     Body(b2.pos, muladd.( b1.m, vmag, b2.v), b2.m))
end#function

@propagate_inbounds update_pos(b::Body, Δt::Float64) =
    Body(muladd.(b.v, Δt, b.pos), b.v, b.m)

function next!(bodies, Δt)
    # Using special iteration tools like eachindex or Iterators.drop
    # is actually measurably slower than for loop around a UnitRange
    @inbounds for i in 1:length(bodies), j in i+1:length(bodies)
        # Script spends roughly 90% of it's time in this loop. Of that
        # 90%, 36% is spent on float multiplication, 34% is spent on
        # the call to muladd, 12% is spent on vector manipulation and
        # iteration overhead, 11% is spent on division, 2% is spent on
        # sqrt, 2% is spent on summing. The remaining 3% is
        # unaccounted for. Overall, multiplication (including muladd)
        # dominates this benchmark.
        bodies[i], bodies[j] = update_velocity(bodies[i], bodies[j], Δt)
    end#for

    # this could be broadcast or mapped, but for-loop is faster
    @inbounds for i in 1:length(bodies)
        bodies[i] = update_pos(bodies[i], Δt)
    end#for
end#function

function energy(bodies)
    e = 0.0
    @inbounds for i in 1:length(bodies)
        bi = bodies[i]
        e += 0.5 * bi.m * sum(bi.v .* bi.v)
        for j in i+1:length(bodies)
            bj = bodies[j]
            Δpos = bi.pos .- bj.pos
            d = √sum(Δpos .* Δpos)
            e = muladd(bi.m * bj.m, -inv(d), e)
        end#for
    end#for
    e
end#function

function main(io, n, Δt=DEFAULT_STEP)
    jupyter = Body(V3d(4.84143144246472090e+00,
                       -1.16032004402742839e+00,
                       -1.03622044471123109e-01),
                   V3d(1.66007664274403694e-03 * DAYS_PER_YEAR,
                       7.69901118419740425e-03 * DAYS_PER_YEAR,
                       -6.90460016972063023e-05 * DAYS_PER_YEAR),
                   9.54791938424326609e-04 * SOLAR_MASS)
    saturn = Body(V3d(8.34336671824457987e+00,
                      4.12479856412430479e+00,
                      -4.03523417114321381e-01),
                  V3d(-2.76742510726862411e-03 * DAYS_PER_YEAR,
                      4.99852801234917238e-03 * DAYS_PER_YEAR,
                      2.30417297573763929e-05 * DAYS_PER_YEAR),
                  2.85885980666130812e-04 * SOLAR_MASS)
    uranus = Body(V3d(1.28943695621391310e+01,
                      -1.51111514016986312e+01,
                      -2.23307578892655734e-01),
                  V3d(2.96460137564761618e-03 * DAYS_PER_YEAR,
                      2.37847173959480950e-03 * DAYS_PER_YEAR,
                      -2.96589568540237556e-05 * DAYS_PER_YEAR),
                  4.36624404335156298e-05 * SOLAR_MASS)
    neptune = Body(V3d(1.53796971148509165e+01,
                       -2.59193146099879641e+01,
                       1.79258772950371181e-01),
                   V3d(2.68067772490389322e-03 * DAYS_PER_YEAR,
                       1.62824170038242295e-03 * DAYS_PER_YEAR,
                       -9.51592254519715870e-05 * DAYS_PER_YEAR),
                   5.15138902046611451e-05 * SOLAR_MASS)
    sun = init_sun((jupyter, saturn, uranus, neptune))
    # Although bodies is fixed size, vector instead of tuple for mutability
    bodies = [sun, jupyter, saturn, uranus, neptune]

    @printf(io, "%.9f\n", energy(bodies))
    for i in 1:n
        next!(bodies, Δt)
    end#for
    @printf(io, "%.9f\n", energy(bodies))
end#function

end#module

isinteractive() || NBodyAB.main(stdout, parse(Int, ARGS[1]))


using Test
using Example2018

@testset "Polynomials form a ring" begin

	z = Polynomial([0])

	for n in 1:200
		deg = rand(0:100)
		p = Polynomial(rand(-100000:100000, deg+1))
		@test p + z == p
		@test z + p == p
		@test p + (-p) == z
	end
end


using Test, LinearMaps, LinearAlgebra, BenchmarkTools

@testset "identity/scaling map" begin
    @test_throws ArgumentError LinearMaps.UniformScalingMap(true, -1)
    A = 2 * rand(ComplexF64, (10, 10)) .- 1
    B = rand(size(A)...)
    M = @inferred 1 * LinearMap(A)
    N = @inferred LinearMap(B)
    LC = @inferred M + N
    v = rand(ComplexF64, 10)
    w = similar(v)
    Id = @inferred LinearMap(I, 10)
    @test occursin("10×10 LinearMaps.UniformScalingMap{Bool}", sprint((t, s) -> show(t, "text/plain", s), Id))
    @test_throws ErrorException LinearMaps.UniformScalingMap(1, 10, 20)
    @test_throws ErrorException LinearMaps.UniformScalingMap(1, (10, 20))
    @test size(Id) == (10, 10)
    @test @inferred isreal(Id)
    @test @inferred issymmetric(Id)
    @test @inferred ishermitian(Id)
    @test @inferred isposdef(Id)
    @test Id * v == v
    @test (2 * M' + 3 * I) * v == 2 * A'v + 3v
    @test (3 * I + 2 * M') * v == 2 * A'v + 3v
    @test (2 * M' - 3 * I) * v == 2 * A'v - 3v
    @test (3 * I - 2 * M') * v == -2 * A'v + 3v
    @test (3 * I - 2 * M') * v == -2 * A'v + 3v
    @test transpose(LinearMap(2 * M' + 3 * I)) * v ≈ transpose(2 * A' + 3 * I) * v
    @test LinearMap(2 * M' + 0I)' * v ≈ (2 * A')' * v
    for λ in (0, 1, rand()), α in (0, 1, rand()), β in (0, 1, rand()), sz in (10, (10,5))
        Λ = @inferred LinearMap(λ*I, 10)
        x = rand(Float64, sz)
        y = rand(Float64, sz)
        b = @benchmarkable mul!($y, $Λ, $x, $α, $β)
        @test run(b, samples=3).allocs == 0
        y = deepcopy(x)
        @inferred mul!(y, Λ, x, α, β)
        @test y ≈ λ * x * α + β * x
    end
    for elty in (Float64, ComplexF64), transform in (identity, transpose, adjoint)
        λ = rand(elty)
        x = rand(10)
        J = @inferred LinearMap(LinearMaps.UniformScalingMap(λ, 10))
        @test transform(J) * x == transform(λ) * x
        J = @inferred LinearMap(λ*I, 10)
        @test (λ * J) * x == (J * λ) * x == (λ * λ) * x
    end
    X = rand(10, 10); Y = similar(X)
    @test mul!(Y, Id, X) == X
    @test Id*X isa LinearMap
    @test X*Id isa LinearMap
    @test Matrix(Id*X) == X
    @test Matrix(X*Id) == X
end


function parse_layout(simple_plot::SimplePlot)
  cur_layout = deepcopy(simple_plot.layout)

  if !isempty(simple_plot.xlims)
    cur_lims = simple_plot.xlims
    if cur_layout["xaxis"]["type"] == "log"
      cur_lims = map(log10, cur_lims)
    end
    cur_layout["xaxis"]["range"] = cur_lims
  end

  if !isempty(simple_plot.ylims)
    cur_lims = simple_plot.ylims
    if cur_layout["yaxis"]["type"] == "log"
      cur_lims = map(log10, cur_lims)
    end
    cur_layout["yaxis"]["range"] = cur_lims
  end

  if !isempty(simple_plot.xticks)
    if length(simple_plot.xticks) == 1
      tickvals = first(simple_plot.xticks)
      ticktext = nothing
    else
      tickvals, ticktext = simple_plot.xticks
    end

    if cur_layout["xaxis"]["type"] == "log"
      tickvals = map(log10, tickvals)
    end

    cur_layout["xaxis"]["tickvals"] = tickvals
    isnothing(ticktext) || ( cur_layout["xaxis"]["ticktext"] = ticktext )
  end

  if !isempty(simple_plot.yticks)
    if length(simple_plot.yticks) == 1
      tickvals = first(simple_plot.yticks)
      ticktext = nothing
    else
      tickvals, ticktext = simple_plot.yticks
    end

    if cur_layout["yaxis"]["type"] == "log"
      tickvals = map(log10, tickvals)
    end

    cur_layout["yaxis"]["tickvals"] = tickvals
    isnothing(ticktext) || ( cur_layout["yaxis"]["ticktext"] = ticktext )
  end

  cur_layout
end


include("../src/PushRelabelMaxFlow.jl")

function clique_expansion_penalties(k)
    """
    Symmetric clique splitting penalty for hyperedge of size k.
    Only gives first r = floor(k/2) penalties
    """
    @assert(k > 1)
    r = floor(Int64,(k)/2)
    I = collect(0:k)
    w = I .* (k .- I)
    w = w[1:r+1]
    return w
end

function clique_expansion_EdgesW(order::Vector{Int64})
    EdgesW = Vector{Vector{Float64}}()
    for k in order
        weight = 1/(k-1)
        push!(EdgesW,weight*clique_expansion_penalties(k))
    end
    return EdgesW
end

function sqrt_EdgesW(order::Vector{Int64})
    EdgesW = Vector{Vector{Float64}}()
    for k in order
        r = floor(Int64,(k)/2)
        w = sqrt.(collect(0:1:r))
        # w = min.(w,delta)
        push!(EdgesW,w)
    end
    return EdgesW
end

function generic_EdgesW(order::Vector{Int64},fun)
    # fun is a generic function
    EdgesW = Vector{Vector{Float64}}()
    for k in order
        r = floor(Int64,(k)/2)
        w = fun(collect(0:1:r))
        push!(EdgesW,w)
    end
    return EdgesW
end

function clique_expansion_EdgesW(Edges::Vector{Vector{Int64}},weighted=true)
    EdgesW = Vector{Vector{Float64}}()
    for e in Edges
        k = length(e)
        if weighted
            weight = 1/k
        else
            weight = 1
        end
        push!(EdgesW,weight*clique_expansion_penalties(k))
    end
    return EdgesW
end

function delta_linear_EdgesW(Edges::Vector{Vector{Int64}},delta::Float64)
    EdgesW = Vector{Vector{Float64}}()
    for e in Edges
        k = length(e)
        r = floor(Int64,(k)/2)
        w = collect(0:1:r)
        w = min.(w,delta)
        push!(EdgesW,w)
    end
    return EdgesW
end

function gen_hypergraph_cut(Edges,EdgesW,S,n)
    """
    For a set S, evaluate generalized hypergraph cut using splitting functions
    from EdgesW (symmetric and cardinality-based).
    """
    eS = zeros(n)
    eS[S] = ones(length(S))

    # higher-order cut penalties
    cutval=0
    for i = 1:length(Edges)
        e = Edges[i]
        k = length(e)
        w = EdgesW[i]
        t = round(Int64,sum(eS[e]))

        # do this because it's a symmetric splitting function,
        # so w is only half the size
        t = min(t,k-t)
        cutval += w[t+1]
    end

    return cutval

end

# For a set S in a graph with adjacency matrix A, return some information about
# S including its conductance, number of interior edges, volume, and cut.
function set_stats(A::SparseMatrixCSC{Float64,Int64},
    S::Vector{Int64},volA::Float64)

    if volA == 0.0
        volA = sum(A.nzval)
    end

    if length(S) == size(A,1)
        # then we have an indicator vector
        S = findall(x->x!=0,eS)
        AS = A[:,S]
    else
        # then we have a subset
        @assert(minimum(S) >= 1)
        @assert(maximum(S) <= size(A,1))
        AS = A[:,S]
    end

    vol = sum(AS.nzval);
    SAS = AS[S,:]
    edges = sum(SAS.nzval);
    cut = vol-edges

    cond = cut/minimum([vol,volA-vol]);

    return cut, vol, edges, cond

end


cd(@__DIR__); include("setups/distROT_distTSD_ratio.jl")
gr(dpi = 200)

## (a) P64
include("setups/path64.jl")
ρ1 = generate_ROT_TSD_ratio(nsim, 𝚽, ∇𝚽, 𝛌, Q)
# plt = ROT_TSD_ratio_histogram(ρ1)
# savefig(plt, "../figs/Path64_ROT_TSD.png")

## (b) P₇ x P₃
include("setups/grid7x3.jl")
ρ2 = generate_ROT_TSD_ratio(nsim, 𝚽, ∇𝚽, 𝛌, Q)
# plt = ROT_TSD_ratio_histogram(ρ2)
# savefig(plt, "../figs/Grid7x3_ROT_TSD.png")

## (c) Erdos Rényi
include("setups/er.jl")
ρ3 = generate_ROT_TSD_ratio(nsim, 𝚽, ∇𝚽, 𝛌, Q)
# plt = ROT_TSD_ratio_histogram(ρ3)
# savefig(plt, "../figs/Erdos_Renyi_ROT_TSD.png")

## (d) weighted RGC100
include("setups/rgc100.jl")
ρ4 = generate_ROT_TSD_ratio(nsim, 𝚽, ∇𝚽, 𝛌, Q; edge_length = edge_length)
# plt = ROT_TSD_ratio_histogram(ρ4)
# savefig(plt, "../figs/wRGC100_ROT_TSD.png")

## boxplot
plt = boxplot(["(a) Path" "(b) Grid" "(c) ER" "(d) RGC100"], [ρ1, ρ2, ρ3, ρ4];
    legend = false, frame = :box, ylim = [0.9, 1.9], tickfontsize = 11,
    outliers = true, grid = false, range = 3, lw = 1, size = (800, 600),
    ylab = "ρ", yguidefontsize = 14, xtickfontsize = 12, color = :white)
savefig(plt, "../figs/ROT_TSD_boxplots.png")

## Table 4.1
display_basic_stats([ρ1, ρ2, ρ3, ρ4])


using ControlSystems
using ControlSystems: AbstractStateSpace, ssdata, ninputs, noutputs

function basis_factor(p::Number)
    (A=p, B=sqrt(1-abs2(p)), C=sqrt(1-abs2(p)), D=-p)
end

# eq (3) in https://people.kth.se/~andersb/research/BF03.pdf
# IDENTIFICATION OF RATIONAL SPECTRAL DENSITIES USING ORTHONORMAL BASIS FUNCTIONS
function expand_basis(l, s) # large, small
    A = [l.A zeros(size(l.A,1), 1); s.B*l.C s.A]
    B = [l.B; s.B*l.D]
    C = [s.D*l.C s.C]
    D = s.D*l.D
    (; A, B, C, D)
end

function kautz(a::AbstractVector)
    h = 1/(pi*maximum(a))
    b = kautz(a, h)
    d2c(b)
end

function kautz(a::AbstractVector, h)
    if maximum(abs, a) > 1
        if all(>(0) ∘ real, a)
            a = -a
        end
        a = exp.(h .* a)
        @assert all(<(1) ∘ abs, a) "Ustable poles does not make sense"
    end
    factors = basis_factor.(a)
    Q = factors[1]
    for f in factors[2:end]
        Q = expand_basis(Q, f)
    end
    T = (Q...,)
    # A,B,C,D = real.(T)
    A,B,C,D = T
    C2 = I(size(A,1))
    ss(A,B,C2,0, h)
end

function laguerre_oo(a::Number, Nq)
    A = diagm(fill(-a, Nq))
    for i = 2:Nq
        A[diagind(A, i-1)] .+= 2a*(-1)^i
    end
    C = sqrt(2a)*(-1).^(0:Nq-1)'
    ss(A,I(Nq),C,0)
end

function laguerre(a, Nq)
    A = diagm(fill(-a, Nq))
    A[diagind(A, 1)] .= 1
    B = diagm(fill(1, Nq))
    for i = 1:Nq-1
        for (k,j) in enumerate(diagind(A, i))
            B[j] = (-1)^i * binomial(i+k-1, k-1)
        end
    end
    Ma = sqrt(2a)*diagm((2a).^(0:Nq-1))
    C = [1 zeros(1, Nq-1)]
    ss(A,Ma*B,C,0)
end

function adhocbasis(a)
    b = map(1:length(a)-1) do i
        ss(leadlinkat(a[i], 4, 1/2)*laglink(a[i+1], 2))
    end
    [b..., leadlinkat(a[end], 4, 1/2)]
end

function simplified_laguerre(a)
    map(enumerate(a)) do (i,ai)
        tf(1, [1.0, ai])^i
    end
end

# Utils ===============================================================================
const SomeBasis = Union{AbstractVector{<:LTISystem}, LTISystem}

basislength(v::AbstractVector) = length(v)
function basislength(v::LTISystem)
    n = ninputs(v)
    n == 1 ? noutputs(v) : n
end

function add_poles(basis::AbstractStateSpace, sys::LTISystem)
    add_poles(basis, ss(sys).A)
end

function add_poles(basis::AbstractStateSpace, Ad)
    A,B,C,D = ssdata(basis)
    T = eltype(A)
    nx,nu,ny = basis.nx,basis.nu,basis.ny
    nn = size(Ad, 1)
    @assert size(A) == size(B) "only supports laguerre_oo basis"
    B == I || @warn("B was not identity, the B matrix will be overwritten with I")

    Ae = [A zeros(size(A,1), nn); zeros(T, nn, nx) Ad]
    # Be = [B; zeros(T, nn, nu)]
    Ce = [C ones(T, ny, nn)]
    ss(Ae,I(nn+nx),Ce,0)
end

function ControlSystems.freqresp(P::LTISystem, basis::SomeBasis, ω::AbstractVector, p)
    @assert ninputs(P) == noutputs(P) == 1 "Only supports SISO"
    resp_P = freqresp(P, ω) |> vec
    basis_resp = basis_responses(basis, ω, inverse=false)
    @assert length(basis_resp) == length(p) "p does not have the same length as the basis"

    FPs = map(basis_resp) do Fi
        Fi .* resp_P
    end

    Gmat = reduce(hcat, FPs)
    Gmat*p
end

function sum_basis(basis::AbstractVector, p::AbstractVector)
    out = ss(p[1]*basis[1])
    for i in 2:length(p)
        out = out + balreal(p[i] * basis[i])[1]
    end
    out
end

function sum_basis(basis::AbstractStateSpace, p::AbstractVector)
    A,B,C,D = ssdata(basis)
    @assert all(iszero, D)
    if ninputs(basis) > 1
        isdiscrete(basis) ? ss(A,B*p,C,0,basis.Ts) : ss(A,B*p,C,0)
    else
        isdiscrete(basis) ? ss(A,B,p'C,0,basis.Ts) : ss(A,B,p'C,0)
    end
end

function basis_responses(basis::AbstractVector, ω; inverse=false)
    inverse && (basis = inv.(basis))
    Gs = freqresp.(basis, Ref(ω))
    vec.(Gs)
end

function basis_responses(basis::AbstractStateSpace, ω; inverse=false)
    inverse && (basis = inv(basis))
    Gs = freqresp(basis, ω)
    if ninputs(basis) > 1
        vec.(eachslice(Gs, dims=3))
    else
        vec.(eachslice(Gs, dims=2))
    end
end

function filter_bank(basis::AbstractStateSpace{<:Discrete}, signal::AbstractVector)
    no = ninputs(basis)
    t = range(0, step=basis.Ts, length=length(signal))
    if no > 1
        y,_,_ = lsim(basis, repeat(signal, 1, no), t)
    else
        y,_,_ = lsim(basis, signal, t)
    end
    y
end

function ωζ2complex(ω, ζ)
    p = ω*cis(pi-acos(ζ))
    p
end

Base.one(::TransferFunction{Continuous, ControlSystems.SisoRational{Float64}}) = tf(1)

reflect(x) = complex(-abs(real(x)), imag(x))
function minimum_phase(G::TransferFunction)
    z,p,k = zpkdata(G) .|> first
    z = reflect.(z)
    p = reflect.(p)
    Gmp = tf(zpk(z,p,k))
    if sign(dcgain(G)[]) != sign(dcgain(Gmp)[])
        Gmp = -Gmp
    end
    Gmp
end

function minimum_phase(G::StateSpace)
    H = tf(G)
    G = ss(minimum_phase(H))
    G = ControlSystems.balance_statespace(G)[1]
    try
        G = balreal(G)[1]
    catch
    end
    G
end


@doc """
	intersect(c1::AbstractCurve,c2::AbstractCurve; tol=<default>)
Find the intersection(s) of two curves. The result could be a vector of zero or more values, or a curve.
""" intersect

# utility, to be used in what follows
function twolines_meet(z1,s1,z2,s2,tol)
	M = [ real(s1) -real(s2); imag(s1) -imag(s2) ]
	if cond(M) > 1/tol
		# parallel lines, caller must decide
		return NaN,NaN
	else
		d = z2 - z1
		t = M \ [real(d);imag(d)]
		return t[1],t[2]
	end
end

function intersect(l1::Line,l2::Line;tol=DEFAULT[:tol])
	s1,s2 = l1.direction,l2.direction
	z1,z2 = l1.base,l2.base
	t1,t2 = twolines_meet(z1,s1,z2,s2,tol)
	if isnan(t1)     # parallel lines
		# either identical or no intersection
		return dist(z2,l1) < tol*(abs(z2)+1) ? l1 : []
	else
		return [z1 + t1*s1]
	end
end

function intersect(g1::Segment,g2::Segment;tol=DEFAULT[:tol])
	z1,z2 = g1.za,g2.za
	s1,s2 = g1.zb-z1,g2.zb-z2
	t1,t2 = twolines_meet(z1,s1,z2,s2,tol)
	if isnan(t1)   # parallel lines
		# scale first segment to [0,1]
		a = real( (g2.za-g1.za)/s1 )
		b = real( (g2.zb-g1.za)/s1 )
		if a > b
			c = b; b = a; a = c
		end
		if b < -tol || a > 1+tol
			return []
		else
			return Segment(z1+max(0,a)*s1,z1+min(1,b)*s1)
		end
	else   # nonparallel
		if (-tol ≤ t1 ≤ 1+tol) && (-tol ≤ t2 ≤ 1+tol)
			return [z1 + t1*s1]
		else
			return []
		end
	end
end

intersect(g::Segment,l::Line;tol=DEFAULT[:tol]) = intersect(l,g,tol=tol)
function intersect(l::Line,g::Segment;tol=DEFAULT[:tol])
	z1,z2 = l.base,g.za
	s1,s2 = l.direction,g.zb-g.za
	t1,t2 = twolines_meet(z1,s1,z2,s2,tol)
	if isnan(t1)   # parallel lines
		return dist(g.za,l) < tol*(1+abs(g.za)) ? g : []
	else
		return 0 ≤ t2 ≤ 1 ? [z1+t1*s1] : []
	end
end

function intersect(r1::Ray,r2::Ray;tol=DEFAULT[:tol])
	s1,s2 = exp(complex(0,r1.angle)),exp(complex(0,r2.angle))
	z1,z2 = r1.base,r2.base
	t1,t2 = twolines_meet(z1,s1,z2,s2,tol)
	if isnan(t1)     # parallel lines
		if dist(z2,r1) < tol*(abs(z2)+1)
			return r2
		else
			dist(z1,r2) < tol*(abs(z1)+1) ? r1 : []
		end
	else
		return (t1 ≥ 0) && (t2 ≥ 0) ? [z1 + t1*s1] : []
	end
end

intersect(r::Ray,l::Line;tol=DEFAULT[:tol]) = intersect(l,r,tol=tol)
function intersect(l::Line,r::Ray;tol=DEFAULT[:tol])
	z1,z2 = l.base,r.base
	s1,s2 = l.direction,exp(1im*r.angle)
	t1,t2 = twolines_meet(z1,s1,z2,s2,tol)
	if isnan(t1)   # parallel lines
		return dist(r.base,l) < tol*(1+abs(r.base)) ? r : []
	else
		return 0 ≤ t2 ? [z1+t1*s1] : []
	end
end

intersect(g::Segment,r::Ray;kw...) = intersect(r,g;kw...)
function intersect(r::Ray,g::Segment;tol=DEFAULT[:tol])
	z1,z2 = r.base,g.za
	s1,s2 = exp(1im*r.angle),g.zb-g.za
	t1,t2 = twolines_meet(z1,s1,z2,s2,tol)
	if isnan(t1)   # parallel lines
		# move ray to positive Re axis
		a = real((g.za-z1)/s1)
		b = real((g.zb-z1)/s1)
		if a > b
			c = b; b = a; a = c
		end
		return b < -tol ? [] : Segment(z1+max(0,a)*s1,z1+b*s1)
	else   # nonparallel
		return (-tol ≤ t1) && (-tol ≤ t2 ≤ 1+tol) ? [z1 + t1*s1]  : []
	end
end

function intersect(c1::Circle,c2::Circle;tol=DEFAULT[:tol])
	r1,r2 = c1.radius,c2.radius
	z1,z2 = c1.center,c2.center
	delta = z2-z1
	d = abs(delta)
	if abs(d) < tol
		if isapprox(r1,r2,rtol=tol,atol=tol)
			return c1
		else
			return []
		end
	elseif (d > r1+r2) || d < abs(r1-r2)
		return []
	else
		a = (r1^2 - r2^2 + d^2)/(2d)
		p = z1 + a*delta/d
		h = sqrt(r1^2-a^2)
		w = 1im*h*delta/d
		z = [p+w,p-w]
		return z
	end
end

intersect(l::Line,c::Circle;kw...) = intersect(c,l;kw...)
function intersect(c::Circle,l::Line;tol=DEFAULT[:tol])
	# find a radius that is perpendicular to the line
	p = Line(c.center,direction=1im*l.direction)
	zi = intersect(p,l)[1]
	a,r = abs(zi-c.center),c.radius
	if a ≤ r + tol
		c = sqrt(r^2-a^2)*l.direction
		if abs(c) > tol
			return [zi+c,zi-c]
		else
			return [zi]
		end
	else
		return []
	end
end

intersect(r::Ray,c::Circle;tol=DEFAULT[:tol]) = intersect(c,r,tol=tol)
function intersect(c::Circle,r::Ray;tol=DEFAULT[:tol])
	# probably not the most efficient way
	z = intersect(c,Line(r.base,direction=exp(1im*r.angle)),tol=tol)
	d = [ dist(z,r) for z in z ]
	return z[d.<=tol]
end

intersect(s::Segment,c::Circle;tol=DEFAULT[:tol]) = intersect(c,s,tol=tol)
function intersect(c::Circle,s::Segment;tol=DEFAULT[:tol])
	# probably not the most efficient way
	z = intersect(c,Line(s.za,s.zb),tol=tol)
	d = [ dist(z,s) for z in z ]
	return z[d.<=tol]
end

function intersect(a1::Arc,a2::Arc;tol=DEFAULT[:tol])
	if isapprox(a1.circle,a2.circle,tol=tol)
		# to what extent do the arcs overlap on the same circle?
		r = a1.circle.radius
		d10ison2 = dist(point(a1,0),a2) < tol*(1+r)
		d11ison2 = dist(point(a1,1),a2) < tol*(1+r)
		d20ison1 = dist(point(a2,0),a1) < tol*(1+r)
		d21ison1 = dist(point(a2,1),a1) < tol*(1+r)
		if d10ison2
			if d11ison2
				return a1
			else
				t = d20ison1 ? arg(a1,point(a2,0)) : arg(a1,point(a2,1))
				return Arc(point(a1,[0,t/2,t]))
			end
		elseif d11ison2
			t = d20ison1 ? arg(a1,point(a2,0)) : arg(a1,point(a2,1))
			return Arc(point(a1,[1-t,1-t/2,1]))
		else
			return []
		end
	else
		z = intersect(a1.circle,a2.circle)
		f = w -> (dist(w,a1) < tol) & (dist(w,a2) < tol)
		return filter(f,z)
	end
end

intersect(c::AbstractCurve,a::Arc;kw...) = intersect(a,c;kw...)
function intersect(a::Arc,c::AbstractCurve;tol=DEFAULT[:tol])
	if isapprox(a.circle,c,tol=tol)
		return a
	else
		z = intersect(a.circle,c)
		return filter(v->dist(v,a)<tol,z)
	end
end

# Determine the (directional) crossings of a Ray(y*1im,0) with a given curve.
# For intersection at an endpoint, give half the value, which gives the right winding number for a polygon.
function raycrossing(z,c::Segment;tol=DEFAULT[:tol])
	s = 0
	y = imag(z)
	a,b = imag(c(0)),imag(c(1))
	atend = abs(a-y) < tol || abs(b-y) < tol
	if y ≥ a
		if y ≤ b && isleft(z,c)
			s = atend ? s+0.5 : s+1
		end
	else
		if y ≥ b && !isleft(z,c)
			s = atend ? s-0.5 : s-1
		end
	end
	return s
end

function raycrossing(z,c::AbstractCurve;tol=DEFAULT[:tol])
	s = 0
	for z in intersect(c,Ray(z,0))
		t = arg(c,z)
		if -tol ≤ t ≤ 1+tol
			y = imag(unittangent(c,t))
			if abs(y) < tol
				# horiztonal tangent; check a perturbed location
				ϵ = sqrt(tol)
				if t < 0.5
					y = imag(unittangent(c,t+2ϵ))
				else
					y = -imag(unittangent(c,t-2ϵ))
				end
			end
			d = tol < t < 1-tol ? copysign(1,y) : copysign(1,y)/2
			s += d
		end
	end
	return s
end


using StatsBase
using Distances

module kNN
    export knn, kernelregression

    using StatsBase
    using Distances
    using NearestNeighbors
	using SmoothingKernels

    include("majority_vote.jl")
    include("classifier.jl")
    include("regress.jl")
end


# SPXL長期模擬データから，運用ケース別の成績を計算する
using CSV, DataFrames, Dates, Plots, StatsBase
pyplot(fmt=:svg, fontfamily="meiryo") # plotの設定

# SPXL読み込み
dfmt = dateformat"yyyy-mm-dd"
spxl = DataFrame(CSV.File("./spxl_pred_1954.csv", dateformat=dfmt, header=true))

# -------------------------------------
# (1)まずはSP500とSPXLで，30年の運用データの中央値，上位下位5%の価格推移を計算する
# -------------------------------------
# SPXLを30年データに分割
# 1954/1/4~1984/1/4, 1954/1/5~1984/1/5...1981/4/28~2021/4/28というイメージ
st_date = spxl.timestamp[1]
data_len = length(spxl.timestamp[spxl.timestamp .<= st_date + Dates.Year(30)])
en_date = spxl.timestamp[data_len]
# データ数data_len(7541日)分のデータを，data_num(9407)個用意
data_num = length(spxl.spxl) - data_len + 1
spxl_30y = ones(data_len, data_num)
for i in 1:data_num
    spxl_30y[:,i] = spxl.spxl[i:i + data_len - 1]
end

# データの正規化(1行目の値を1とした比率に変換)
for c = 1:data_num
    spxl_30y[:,c] ./= spxl_30y[1,c]
end

# 5パーセンタイル，50パーセンタイル, 95パーセンタイルを計算してプロット
quan_spxl = zeros(data_len, 3)
for r in 1:data_len
    quan_spxl[r,:] = quantile(spxl_30y[r,:], [0.05, 0.5, 0.95])
end

labels = ["下位5%", "中央値", "上位5%", "投資額"]
plot( (1:data_len) .* 30 / data_len, quan_spxl[:,3], label=labels[3], linecolor=:blue, xlabel="年", ylabel="価格[\$]", yscale=:log10, ylims=(10^-0.5, 10^2.5))
plot!( (1:data_len) .* 30 / data_len, quan_spxl[:,2], label=labels[2], linecolor=:green, ms=1, yscale=:log10)
plot!( (1:data_len) .* 30 / data_len, quan_spxl[:,1], label=labels[1], linecolor=:red, yscale=:log10)
savefig("SPXL_30y_5parcentile.png")

# sp500を30年データに分割--------------------
sp500_30y = ones(data_len, data_num)
for i in 1:data_num
    sp500_30y[:,i] = spxl.sp500_yidx[i:i + data_len - 1]
end

# データの正規化(1行目の値を1とした比率に変換)
for c = 1:data_num
    sp500_30y[:,c] ./= sp500_30y[1,c]
end

# 5パーセンタイル，50パーセンタイル, 95パーセンタイルを計算してプロット
quan_sp500 = zeros(data_len, 3)
for r in 1:data_len
    quan_sp500[r,:] = quantile(sp500_30y[r,:], [0.05, 0.5, 0.95])
end

plot( (1:data_len) .* 30 / data_len, quan_sp500[:,3], label=labels[3], linecolor=:blue, xlabel="年", ylabel="価格[\$]", yscale=:log10, ylims=(10^-0.5, 10^2.5))
plot!( (1:data_len) .* 30 / data_len, quan_sp500[:,2], label=labels[2], linecolor=:green, yscale=:log10)
plot!( (1:data_len) .* 30 / data_len, quan_sp500[:,1], label=labels[1], linecolor=:red, yscale=:log10)
savefig("SP500_30y_5parcentile.png")

# -------------------------------------
# (2)上記の価格推移で，500万+毎月10万を積立投資した場合を計算する
# -------------------------------------

# 年月日から購入金額を計算
purchases = zeros(data_len)
purchases[1] = 500 # 初回500万円
mid_flag = false
for i = 2:data_len
    if (Dates.Month(spxl.timestamp[i - 1]) != Dates.Month(spxl.timestamp[i]))
        purchases[i] = 10
        mid_flag = false
    elseif ( (Dates.Month(spxl.timestamp[i]) == Dates.Month(6) || Dates.Month(spxl.timestamp[i]) == Dates.Month(12)) &&
        Dates.Day(spxl.timestamp[i]) > Dates.Day(15) && mid_flag == false)
        purchases[i] = 20
        mid_flag = true
    end
end

# 購入金額を積算して，総投資額も計算しておく
total_purchases = accumulate(+, purchases)

# 購入株式数
quantity_spxl = zeros(size(spxl_30y))
for n in 1:data_num
    quantity_spxl[:,n] = purchases ./ spxl_30y[:,n]
end

# 購入株式数を積算して保有株式数とする
stocks_spxl = zeros(size(spxl_30y))
for n in 1:data_num
    stocks_spxl[:,n] = accumulate(+, quantity_spxl[:,n])
end

# 保有株式数にその時の価格を掛け算して，資産額を算出
assets_spxl = copy(stocks_spxl)
for c = 1:data_num
    assets_spxl[:,c] .*= spxl_30y[:,c]
end

# 5パーセンタイル，50パーセンタイル, 95パーセンタイルを計算してプロット
quan_spxl_assets = zeros(data_len, 3)
for r in 1:data_len
    quan_spxl_assets[r,:] = quantile(assets_spxl[r,:], [0.05, 0.5, 0.95])
end

plot( (1:data_len) .* 30 / data_len, quan_spxl_assets[:,3], label=labels[3], linecolor=:blue, xlabel="年", ylabel="資産額[万円]", yscale=:log10, ylims=(10^2.5, 10^6.0))
plot!( (1:data_len) .* 30 / data_len, quan_spxl_assets[:,2], label=labels[2], linecolor=:green, yscale=:log10)
plot!( (1:data_len) .* 30 / data_len, quan_spxl_assets[:,1], label=labels[1], linecolor=:red, yscale=:log10)
plot!( (1:data_len) .* 30 / data_len, total_purchases, label=labels[4], linecolor=:black, yscale=:log10)
savefig("SPXL_30y_5parcentile_assets.png")

# 成績の計算
for i = 3:-1:1
    print(labels[i] * ": " * string(quan_spxl_assets[end, i] * (1 - 0.20315)) * "万円\n")
end
quan_spxl_assets[2513, 2]

print("下位5%の最大下落率[%]:" * string(100 - minimum(quan_spxl_assets[:,1] ./ total_purchases) * 100)[1:4])

# SP500も計算-------------------------
# 購入株式数
quantity_sp500 = zeros(size(sp500_30y))
for n in 1:data_num
    quantity_sp500[:,n] = purchases ./ sp500_30y[:,n]
end

# 購入株式数を積算して保有株式数とする
stocks_sp500 = zeros(size(sp500_30y))
for n in 1:data_num
    stocks_sp500[:,n] = accumulate(+, quantity_sp500[:,n])
end

# 保有株式数にその時の価格を掛け算して，資産額を算出
assets_sp500 = copy(stocks_sp500)
for c = 1:data_num
    assets_sp500[:,c] .*= sp500_30y[:,c]
end

# 5パーセンタイル，50パーセンタイル, 95パーセンタイルを計算してプロット
quan_sp500_assets = zeros(data_len, 3)
for r in 1:data_len
    quan_sp500_assets[r,:] = quantile(assets_sp500[r,:], [0.05, 0.5, 0.95])
end

plot( (1:data_len) .* 30 / data_len, quan_sp500_assets[:,3], label=labels[3], linecolor=:blue, xlabel="年", ylabel="資産額[万円]", yscale=:log10, ylims=(10^2.5, 10^6.0))
plot!( (1:data_len) .* 30 / data_len, quan_sp500_assets[:,2], label=labels[2], linecolor=:green, yscale=:log10)
plot!( (1:data_len) .* 30 / data_len, quan_sp500_assets[:,1], label=labels[1], linecolor=:red, yscale=:log10)
plot!( (1:data_len) .* 30 / data_len, total_purchases, label=labels[4], linecolor=:black, yscale=:log10)
savefig("SP500_30y_5parcentile_assets.png")
# 成績の計算
for i = 3:-1:1
    print(labels[i] * ": " * string(quan_sp500_assets[end, i] * (1 - 0.20315))[1:6] * "万円\n")
end
quan_sp500_assets[2513, 2] # 10年時点の中央値

print("下位5%の最大下落率[%]:" * string(100 - minimum(quan_sp500_assets[:,1] ./ total_purchases) * 100)[1:4])

# -------------------------------------
# (3)運用パターン変更(完全積立投資)
# -------------------------------------
# 5300万を全て積立投資する．30*12=360ヶ月あるので，14.7万円ずつ．

# 年月日から購入金額を計算
purchases = zeros(data_len)
amount = (5300 / 360)
purchases[1] = amount # 初回
for i = 2:data_len
    if (Dates.Month(spxl.timestamp[i - 1]) != Dates.Month(spxl.timestamp[i])) # 月が変わったタイミングで投資
        purchases[i] = amount
    end
end

# 購入金額を積算して，総投資額も計算しておく
total_purchases = accumulate(+, purchases)

# 購入株式数
quantity_spxl = zeros(size(spxl_30y))
for n in 1:data_num
    quantity_spxl[:,n] = purchases ./ spxl_30y[:,n]
end

# 購入株式数を積算して保有株式数とする
stocks_spxl = zeros(size(spxl_30y))
for n in 1:data_num
    stocks_spxl[:,n] = accumulate(+, quantity_spxl[:,n])
end

# 保有株式数にその時の価格を掛け算して，資産額を算出
assets_spxl = copy(stocks_spxl)
for c = 1:data_num
    assets_spxl[:,c] .*= spxl_30y[:,c]
end

# 5パーセンタイル，50パーセンタイル, 95パーセンタイルを計算してプロット
quan_spxl_assets = zeros(data_len, 3)
for r in 1:data_len
    quan_spxl_assets[r,:] = quantile(assets_spxl[r,:], [0.05, 0.5, 0.95])
end

plot( (1:data_len) .* 30 / data_len, quan_spxl_assets[:,3], label=labels[3], linecolor=:blue, xlabel="年", ylabel="資産額[万円]", yscale=:log10)
plot!( (1:data_len) .* 30 / data_len, quan_spxl_assets[:,2], label=labels[2], linecolor=:green, yscale=:log10)
plot!( (1:data_len) .* 30 / data_len, quan_spxl_assets[:,1], label=labels[1], linecolor=:red, yscale=:log10)
plot!( (1:data_len) .* 30 / data_len, total_purchases, label=labels[4], linecolor=:black, yscale=:log10)
savefig("SPXL_30y_5parcentile_assets_積立時.png")

# 成績の計算
for i = 3:-1:1
    print(labels[i] * ": " * string(quan_spxl_assets[end, i] * (1 - 0.20315))[1:7] * "万円\n")
end
quan_spxl_assets[2513, 2] # 10年時点の中央値

print("下位5%の最大下落率[%]:" * string(100 - minimum(quan_spxl_assets[:,1] ./ total_purchases) * 100)[1:4])

# SP500も計算-------------------------
# 購入株式数
quantity_sp500 = zeros(size(sp500_30y))
for n in 1:data_num
    quantity_sp500[:,n] = purchases ./ sp500_30y[:,n]
end

# 購入株式数を積算して保有株式数とする
stocks_sp500 = zeros(size(sp500_30y))
for n in 1:data_num
    stocks_sp500[:,n] = accumulate(+, quantity_sp500[:,n])
end

# 保有株式数にその時の価格を掛け算して，資産額を算出
assets_sp500 = copy(stocks_sp500)
for c = 1:data_num
    assets_sp500[:,c] .*= sp500_30y[:,c]
end

# 5パーセンタイル，50パーセンタイル, 95パーセンタイルを計算してプロット
quan_sp500_assets = zeros(data_len, 3)
for r in 1:data_len
    quan_sp500_assets[r,:] = quantile(assets_sp500[r,:], [0.05, 0.5, 0.95])
end

plot( (1:data_len) .* 30 / data_len, quan_sp500_assets[:,3], label=labels[3], linecolor=:blue, xlabel="年", ylabel="資産額[万円]", yscale=:log10)
plot!( (1:data_len) .* 30 / data_len, quan_sp500_assets[:,2], label=labels[2], linecolor=:green, yscale=:log10)
plot!( (1:data_len) .* 30 / data_len, quan_sp500_assets[:,1], label=labels[1], linecolor=:red, yscale=:log10)
plot!( (1:data_len) .* 30 / data_len, total_purchases, label=labels[4], linecolor=:black, yscale=:log10)
savefig("SP500_30y_5parcentile_assets_積立時.png")

# 成績の計算
for i = 3:-1:1
    print(labels[i] * ": " * string(quan_sp500_assets[end, i] * (1 - 0.20315))[1:7] * "万円\n")
end
quan_sp500_assets[2513, 2] # 10年時点の中央値

print("下位5%の最大下落率[%]:" * string(100 - minimum(quan_sp500_assets[:,1] ./ total_purchases) * 100)[1:4])

# -------------------------------------
# (4)現金-SPXLリバランス
# -------------------------------------
# 最初に500万円投資，30年間現金は保有，50:50 or 70:30で毎月頭にリバランス
# ただし，平均購入金額を下回ったら売却はしない．

# 年月日から投資への現金追加額を計算
purchases = zeros(data_len)
purchases[1] = 500 # 初回500万円
mid_flag = false
for i = 2:data_len
    if (Dates.Month(spxl.timestamp[i - 1]) != Dates.Month(spxl.timestamp[i]))
        purchases[i] = 10
        mid_flag = false
    elseif ( (Dates.Month(spxl.timestamp[i]) == Dates.Month(6) || Dates.Month(spxl.timestamp[i]) == Dates.Month(12)) &&
        Dates.Day(spxl.timestamp[i]) > Dates.Day(15) && mid_flag == false)
        purchases[i] = 20
        mid_flag = true
    end
end
# 購入金額を積算して，総投資額も計算しておく
total_purchases = accumulate(+, purchases)

# リバランス計算を関数化
# price: 株価推移
# ratio: spxl対現金をいくつにするか．0.7なら70:30
# purchases: 追加投資現金
function rebalance(prices, ratio, purchases, timestamp, fee)
    # for debug
    # prices = spxl_30y[:,1]
    # ratio = 0.7
    # timestamp = spxl.timestamp
    # fee = 0.45/100
    # 変数定義と初期化
    tax_rate = 20.315 / 100 # 税率
    ave_unit_price = prices[1] # 平均取得単価
    hold_stock = zeros(length(prices)) # 保有株式数推移
    hold_stock[1] = purchases[1] * ratio / prices[1]
    hold_cash = zeros(length(prices)) # 保有現金推移
    hold_cash[1] = purchases[1] * (1 - ratio)
    ratio_temp = zeros(length(prices)) # リバランス時の比率
    ratio_temp[1] = ratio
    # 2日目以降を計算
    for i in 2:length(prices)
        # 現金を追加
        hold_cash[i] = purchases[i]
        # 月初めにリバランス
        if (Dates.Month(timestamp[i - 1]) != Dates.Month(timestamp[i]))
            # 現金と株式の比率を計算
            hold_stock_price = sum(hold_stock) * prices[i] # 現在の株式評価額
            hold_cash_price = sum(hold_cash) # 手持ちの現金
            # ratio_temp[i] = hold_stock_price / (hold_cash_price + hold_stock_price) # 現状の株式：現金比率
            diff = ratio * (hold_cash_price + hold_stock_price) - hold_stock_price # 理想の株式価格までの差額
            # 差額が正なら追加投資
            if (diff > 0)
                hold_stock[i] += diff / prices[i] # 差額分の数量の株を追加購入
                hold_cash[i] -= diff * (1 + fee) # 購入+手数料分の現金を引く
                ave_unit_price = ( (sum(hold_stock[1:i - 1]) * ave_unit_price) + diff ) / sum(hold_stock) # 平均取得単価の更新
                # print(ave_unit_price)
                # print("\n")
            # 差額が負なら売却
            else
                # 株価が平均取得単価以上のときだけ売却．下回っていたらホールド．税金がかかるので，その分考慮して多めに売却する
                if (prices[i] >= ave_unit_price)
                    sale = -diff # 売却額
                    profit = sale - sale / prices[i] * ave_unit_price
                    tax = profit * tax_rate # 税金
                    hold_stock[i] -= sale / prices[i] # 差額分の数量の株を売却
                    hold_cash[i] += sale * (1 - fee) - tax # 売却分(税額・手数料考慮)の現金を追加
                end
            end
            # 現金と株式の比率を計算
            hold_stock_price = sum(hold_stock) * prices[i] # 現在の株式評価額
            hold_cash_price = sum(hold_cash) # 手持ちの現金
            ratio_temp[i] = hold_stock_price / (hold_cash_price + hold_stock_price) # 現状の株式：現金比率
        end
    end

    return (hold_stock, hold_cash)
end

hold_stocks = zeros(size(spxl_30y))
hold_cash = zeros(size(spxl_30y))
for n in 1:data_num
    prices = spxl_30y[:,n]
    ratio = 0.7
    timestamp = spxl.timestamp
    fee = 0.45 / 100
    (stock, cash) = rebalance(prices, ratio, purchases, timestamp, fee)
    hold_stocks[:,n] = stock
    hold_cash[:,n] = cash
end

# 購入株式数を積算して保有株式数とする
stocks_spxl = zeros(size(spxl_30y))
for n in 1:data_num
    stocks_spxl[:,n] = accumulate(+, hold_stocks[:,n])
end

# 保有株式数にその時の価格を掛け算して，資産額を算出
assets_spxl = copy(stocks_spxl)
for c = 1:data_num
    assets_spxl[:,c] .*= spxl_30y[:,c]
end

# 現金も積算する
total_cash = zeros(size(spxl_30y))
for n in 1:data_num
    total_cash[:,n] = accumulate(+, hold_cash[:,n])
end

# 5パーセンタイル，50パーセンタイル, 95パーセンタイルを計算してプロット
quan_assets = zeros(data_len, 3)
for r in 1:data_len
    quan_assets[r,:] = quantile(assets_spxl[r,:], [0.05, 0.5, 0.95])
end
quan_cash = zeros(data_len, 3)
for r in 1:data_len
    quan_cash[r,:] = quantile(total_cash[r,:], [0.05, 0.5, 0.95])
end

plot( (1:data_len) .* 30 / data_len, quan_assets[:,3], label="金融資産額" * labels[3], linecolor=:blue, xlabel="年", ylabel="資産額[万円]", yscale=:log10)
plot!( (1:data_len) .* 30 / data_len, quan_assets[:,2], label="金融資産額" * labels[2], linecolor=:green)
plot!( (1:data_len) .* 30 / data_len, quan_assets[:,1], label="金融資産額" * labels[1], linecolor=:red)
plot!( (1:data_len) .* 30 / data_len, quan_cash[:,3], label="現金" * labels[3], linestyle=:dot, linecolor=:blue)
plot!( (1:data_len) .* 30 / data_len, quan_cash[:,2], label="現金" * labels[2], linestyle=:dot, linecolor=:green)
plot!( (1:data_len) .* 30 / data_len, quan_cash[:,1], label="現金" * labels[1], linestyle=:dot, linecolor=:red)
plot!( (1:data_len) .* 30 / data_len, quan_assets[:,3] .+ quan_cash[:,3], label="合計" * labels[3], linestyle=:dash, linecolor=:blue)
plot!( (1:data_len) .* 30 / data_len, quan_assets[:,2] .+ quan_cash[:,2], label="合計" * labels[2], linestyle=:dash, linecolor=:green)
plot!( (1:data_len) .* 30 / data_len, quan_assets[:,1] .+ quan_cash[:,1], label="合計" * labels[1], linestyle=:dash, linecolor=:red)
plot!( (1:data_len) .* 30 / data_len, total_purchases, label=labels[4], linecolor=:black)
savefig("SPXL_30y_5parcentile_assets_70_30リバランス時.png")

# 成績の計算
for i = 3:-1:1
    print(labels[i] * ": " * string((quan_assets[end, i] * (1 - 0.20315) + quan_cash[end,i]))[1:7] * "万円\n")
end
quan_assets[2513, 2] + quan_cash[2513, 2] # 10年時点の中央値

print("下位5%の最大下落率[%]:" * string(100 - minimum((quan_assets[:,1] .+ quan_cash[:,1]) ./ total_purchases) * 100)[1:4])

# -------------------------------------
# (5)15年かけて購入，15年かけて売却
# -------------------------------------
# 年月日から購入金額を計算
data_len_2 = Int(floor(data_len / 2))
purchases = zeros(data_len)
amount = (5300 / (15 * 12))
purchases[1] = amount # 初回も同額
# 15年目まで購入
for i = 2:data_len_2
    if (Dates.Month(spxl.timestamp[i - 1]) != Dates.Month(spxl.timestamp[i])) # 月が変わったタイミングで投資
        purchases[i] = amount
    end
end

# 購入金額を積算して，総投資額も計算しておく
total_purchases = accumulate(+, purchases)

# 購入株式数を購入金額から計算
buy_spxl = zeros(size(spxl_30y))
for n in 1:data_num
    buy_spxl[:,n] = purchases ./ spxl_30y[:,n]
end

# 売却株式数(きっかり15年で同量の株を売却する)
sale_spxl = zeros(size(spxl_30y))
for n in 1:data_num
    sale_spxl[data_len_2:data_len_2 * 2 - 1,n] = buy_spxl[1:data_len_2,n]
end

# 購入株式数と売却株式数を積算して保有株式数を計算する
total_holdings = copy(buy_spxl)
for c = 1:data_num
    total_buy = accumulate(+, buy_spxl[:,c]) # 購入株式数の積算
    total_sale = accumulate(+, sale_spxl[:,c]) # 売却株式数の積算
    total_holdings[:,c] = total_buy .- total_sale # 保有株式数
end
# 保有株式数にその時の価格を掛け算して，金融資産額を計算
holdings_price = copy(total_holdings)
for c = 1:data_num
    holdings_price[:,c] .*= spxl_30y[:,c]
end

# 売却株式数にその時の価格を掛け算して，売却金額を算出
sale_price_spxl = copy(sale_spxl)
for c = 1:data_num
    sale_price_spxl[:,c] .*= spxl_30y[:,c]
end

# 売却額を積算して現金保有額を計算
cash_spxl = zeros(size(spxl_30y))
for n in 1:data_num
    cash_spxl[:,n] = accumulate(+, sale_price_spxl[:,n])
end

# 5パーセンタイル，50パーセンタイル, 95パーセンタイルを計算してプロット
quan_spxl_holdings = zeros(data_len, 3)
for r in 1:data_len
    quan_spxl_holdings[r,:] = quantile(holdings_price[r,:], [0.05, 0.5, 0.95])
end
quan_spxl_cash = zeros(data_len, 3)
for r in 1:data_len
    quan_spxl_cash[r,:] = quantile(cash_spxl[r,:], [0.05, 0.5, 0.95])
end

# 対数グラフだと0が表示できないので，値を微妙に補正
quan_spxl_holdings .+= 10
quan_spxl_cash .+= 10

plot( (1:data_len) .* 30 / data_len, quan_spxl_holdings[:,3], label="金融資産額" * labels[3], linecolor=:blue, xlabel="年", ylabel="資産額[万円]", yscale=:log10)
plot!( (1:data_len) .* 30 / data_len, quan_spxl_holdings[:,2], label="金融資産額" * labels[2], linecolor=:green)
plot!( (1:data_len) .* 30 / data_len, quan_spxl_holdings[:,1], label="金融資産額" * labels[1], linecolor=:red)
plot!( (1:data_len) .* 30 / data_len, quan_spxl_cash[:,3], label="売却合計額" * labels[3], linestyle=:dot, linecolor=:blue)
plot!( (1:data_len) .* 30 / data_len, quan_spxl_cash[:,2], label="売却合計額" * labels[2], linestyle=:dot, linecolor=:green)
plot!( (1:data_len) .* 30 / data_len, quan_spxl_cash[:,1], label="売却合計額" * labels[1], linestyle=:dot, linecolor=:red)
plot!( (1:data_len) .* 30 / data_len, total_purchases, label=labels[4], linecolor=:black)
savefig("SPXL_30y_5parcentile_assets_15年積立15年売却時.png")

# 成績の計算
for i = 3:-1:1
    print(labels[i] * ": " * string(quan_spxl_cash[end, i] * (1 - 0.20315) / total_purchases[end])[1:4] * "倍\n")
end
print("下位5%の最大下落率[%]:" * string(100 - minimum((quan_spxl_holdings[:,1] ./ total_purchases)[1:data_len_2]) * 100)[1:4])

# -------------------------------------
# (6)タイミング投資(直近高値より10%下落で購入)
# -------------------------------------
# 最初に200万円投資，残り300万あり．毎月10万＋ボーナス月は＋20万追加．
# 直近高値に対して10%を下回っていた場合，下落幅に対して以下の式で追加投資を毎日する．
# ((下落率-8)^1.1)+12)*0.2[万円]

# 年月日から投資への現金追加額を計算
purchases = zeros(data_len)
purchases[1] = 500 # 初回500万円
mid_flag = false
for i = 2:data_len
    if (Dates.Month(spxl.timestamp[i - 1]) != Dates.Month(spxl.timestamp[i]))
        purchases[i] = 10
        mid_flag = false
    elseif ( (Dates.Month(spxl.timestamp[i]) == Dates.Month(6) || Dates.Month(spxl.timestamp[i]) == Dates.Month(12)) &&
        Dates.Day(spxl.timestamp[i]) > Dates.Day(15) && mid_flag == false)
        purchases[i] = 20
        mid_flag = true
    end
end
# 購入金額を積算して，総投資額も計算しておく
total_purchases = accumulate(+, purchases)

# ナンピン投資を関数化
# price: 株価推移
# purchases: 追加投資現金
function additional(prices, purchases, timestamp, fee)
    # for debug
    # prices = spxl_30y[:,1]
    # timestamp = spxl.timestamp
    # fee = 0.45/100

    # 変数定義と初期化
    recent_high_price = prices[1] # 直近高値
    hold_stock = zeros(length(prices)) # 保有株式数推移
    hold_stock[1] = purchases[1] * 2.0 / 5.0 / prices[1]
    hold_cash = zeros(length(prices)) # 保有現金推移
    hold_cash[1] = purchases[1] * 3.0 / 5.0
    decrease_rate_total = zeros(length(prices)) # 下落幅

    # 2日目以降を計算
    for i in 2:length(prices)
        # 保有現金を更新
        hold_cash[i] = hold_cash[i - 1] + purchases[i]
        # 直近高値を更新=直近1ヶ月の最高値
        recent = 10
        if (i <= recent)
            recent_high_price = maximum(prices[1:i])
        else
            recent_high_price = maximum(prices[i - recent:i])
        end

        if (hold_cash[i] > 0)
            # 高値から10%を下回っていたら
            decrease_rate = 100.0 - prices[i] / recent_high_price * 100.0 # 下落率
            decrease_rate_total[i] = decrease_rate
            if ( decrease_rate > 10.0 )
                # 投資額を計算して，その分の株式を購入
                amount = (((decrease_rate - 8)^1.1) + 12) * 0.2
                if ( amount > hold_cash[i] ) # 保有現金が不足した場合，保有現金分しか買わない
                    amount = hold_cash[i] / (1 + fee)
                end
                hold_stock[i] = amount / prices[i]
                hold_cash[i] -= amount * (1 + fee)
            end
        end
    end

    return (hold_stock, hold_cash)
end

hold_stocks = zeros(size(spxl_30y))
hold_cash = zeros(size(spxl_30y))
for n in 1:data_num
    prices = spxl_30y[:,n]
    timestamp = spxl.timestamp
    fee = 0.45 / 100
    (stock, cash) = additional(prices, purchases, timestamp, fee)
    hold_stocks[:,n] = stock
    hold_cash[:,n] = cash
end

# 購入株式数を積算して保有株式数とする
stocks_spxl = zeros(size(spxl_30y))
for n in 1:data_num
    stocks_spxl[:,n] = accumulate(+, hold_stocks[:,n])
end

# 保有株式数にその時の価格を掛け算して，資産額を算出
assets_spxl = copy(stocks_spxl)
for c = 1:data_num
    assets_spxl[:,c] .*= spxl_30y[:,c]
end

# 5パーセンタイル，50パーセンタイル, 95パーセンタイルを計算してプロット
quan_assets = zeros(data_len, 3)
for r in 1:data_len
    quan_assets[r,:] = quantile(assets_spxl[r,:], [0.05, 0.5, 0.95])
end
quan_cash = zeros(data_len, 3)
for r in 1:data_len
    quan_cash[r,:] = quantile(hold_cash[r,:], [0.05, 0.5, 0.95])
end
# quan_cashが0だと対数グラフで描画できないので1万円足す
quan_cash .+= 1

#=
plot( (1:data_len) .* 30 / data_len, quan_assets[:,3], label="金融資産額" * labels[3], linecolor=:blue, xlabel="年", ylabel="資産額[万円]", yscale=:log10)
plot!( (1:data_len) .* 30 / data_len, quan_assets[:,2], label="金融資産額" * labels[2], linecolor=:green)
plot!( (1:data_len) .* 30 / data_len, quan_assets[:,1], label="金融資産額" * labels[1], linecolor=:red)
plot!( (1:data_len) .* 30 / data_len, quan_cash[:,3], label="現金" * labels[3], linestyle=:dot, linecolor=:blue)
plot!( (1:data_len) .* 30 / data_len, quan_cash[:,2], label="現金" * labels[2], linestyle=:dot, linecolor=:green)
plot!( (1:data_len) .* 30 / data_len, quan_cash[:,1], label="現金" * labels[1], linestyle=:dot, linecolor=:red) =#
plot( (1:data_len) .* 30 / data_len, quan_assets[:,3] .+ quan_cash[:,3], label="合計" * labels[3], linecolor=:blue, yscale=:log10)
plot!( (1:data_len) .* 30 / data_len, quan_assets[:,2] .+ quan_cash[:,2], label="合計" * labels[2], linecolor=:green)
plot!( (1:data_len) .* 30 / data_len, quan_assets[:,1] .+ quan_cash[:,1], label="合計" * labels[1], linecolor=:red)
plot!( (1:data_len) .* 30 / data_len, total_purchases, label=labels[4], linecolor=:black)
savefig("SPXL_30y_5parcentile_assets_タイミング投資.png")

# 成績の計算
for i = 3:-1:1
    print(labels[i] * ": " * string((quan_assets[end, i] * (1 - 0.20315) + quan_cash[end,i]) / total_purchases[end])[1:4] * "倍\n")
end
(quan_assets[2513, 2] + quan_cash[2513, 2]) / total_purchases[2513]
for i = 3:-1:1
    print(labels[i] * ": " * string((quan_assets[end, i] * (1 - 0.20315) + quan_cash[end,i]))[1:7] * "万円\n")
end
quan_assets[2513, 2] + quan_cash[2513, 2]

print("下位5%の最大下落率[%]:" * string(100 - minimum((quan_assets[:,1] .+ quan_cash[:,1]) ./ total_purchases) * 100)[1:4])