"""
    ConstantLinearAdvection{T,Dim}

Linear advection equation with constant coefficients
``
  \\partial_t u + \\sum_{i=1}^{\\mathrm{Dim}} \\partial_{x_i} u = 0
``
in `Dim` space dimensions using `T` as scalar type.
"""
struct ConstantLinearAdvection{T,Dim} <: ScalarBalanceLaw{T,Dim} end

function ConstantLinearAdvection(T=Float64, Dim=1)
  ConstantLinearAdvection{T,Dim}()
end

function Base.show(io::IO, model::ConstantLinearAdvection{T,Dim}) where {T,Dim}
  print(io, "Linear advection equation with constant coefficient {T=", T, ", Dim=", Dim, "}",
            " with flux f(u) = u * (1,...,1)")
end

"""
    flux{T}(u, model::ConstantLinearAdvection{T,1})

Compute the flux of `u` for `model`.
"""
@inline flux(u, model::ConstantLinearAdvection{T,1}) where {T} = u

"""
    speed(u::Real, model::ConstantLinearAdvection)

Compute the speed f'(`u`) for `model`.
"""
@inline speed(u, model::ConstantLinearAdvection) = one(u)

@inline flux(u, model::ConstantLinearAdvection{T,2}, direction) where {T} = u

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

"""
    (::GodunovFlux){T}(uₗ::T, uᵣ::T, model::ConstantLinearAdvection{T,1})

Compute Godunov's flux between `uₗ` and `uᵣ` for `model`.
"""
function (::GodunovFlux)(uₗ::T, uᵣ::T, model::ConstantLinearAdvection{T,1}) where {T}
    uₗ
end

function (flux::EnergyConservativeFlux)(uₗ::T, uᵣ::T, model::ConstantLinearAdvection{T,1}) where T<:Real
    (uₗ + uᵣ) / 2
end

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

@inline function add_numerical_fluxes_inner_loop2!(du, fluxes, u, balance_law::ConstantLinearAdvection,
                                                    Nx, basis, boundaries_included::Val{false},
                                                    jacx, parallel)
    Pp1 = length(basis.nodes)
    Rl = basis.interp_left
    Rr = basis.interp_right

    # add numerical fluxes
    @inbounds for ix in Base.OneTo(Nx)
        Rul = zero(eltype(u))
        Rur = zero(eltype(u))
        for nx in Base.OneTo(Pp1)
            Rul += Rl[nx]*u[nx,ix]
            Rur += Rr[nx]*u[nx,ix]
        end

        for nx in Base.OneTo(Pp1)
            du[nx,ix] += ((fluxes[ix] - Rul) * Rl[nx] - (fluxes[ix+1] - Rur) * Rr[nx]
                            ) * jacx / basis.weights[nx]
        end
    end

    nothing
end

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

"""
    ConstantLinearAdvectionRiemannSolution{T,T1}

The solution of a Riemann problem `prob` for the linear advection equation with
constant coefficients.
"""
struct ConstantLinearAdvectionRiemannSolution{T,T1} <: ScalarRiemannSolution
    prob::RiemannProblem{ConstantLinearAdvection{T,1},T,T1}
    σ::T
end

"""
    minmax_speeds(sol::ConstantLinearAdvectionRiemannSolution)

Return the minimal and maximal speeds `σ⁻, σ⁺` that appear in the solution `sol`.
"""
function minmax_speeds(sol::ConstantLinearAdvectionRiemannSolution)
    sol.σ, sol.σ
end

"""
    (sol::ConstantLinearAdvectionRiemannSolution)(ξ::Real)

Evaluate the solution `sol` at the value `ξ` of the self-similarity variable
`ξ = (x - x₀) / (t - t₀)`.
"""
function (sol::ConstantLinearAdvectionRiemannSolution)(ξ::Real)
    @unpack σ = sol
    @unpack uₗ, uᵣ = sol.prob

    if ξ < σ
        uₗ
    else
        uᵣ
    end
end

"""
    (sol::ConstantLinearAdvectionRiemannSolution)(t::Real, x::Real)

Evaluate the solution `sol` at the time and space coordinates `t` and `x`.
"""
function (sol::ConstantLinearAdvectionRiemannSolution)(t::Real, x::Real)
    @unpack x₀, t₀ = sol.prob

    sol((x-x₀)/(t-t₀))
end

"""
    solve{T,T1}(prob::RiemannProblem{ConstantLinearAdvection{T,1},T,T1})

Compute the solution of the Riemann prolem `prob`.
"""
function solve(prob::RiemannProblem{ConstantLinearAdvection{T,1},T,T1}) where {T,T1}
    ConstantLinearAdvectionRiemannSolution(prob, one(T))
end


abstract type AbstractNonLinearLeastSquares end

mutable struct NonLinearLeastSquares{BFVT, BFJVT} <: AbstractNonLinearLeastSquares
    # Basis functions where f[i] ∀ i ∈ [1, k] should be a function
    # which takes the state and independant
    # variable as arguments (i.e f[i](x,tₖ) where tₖ is the time at
    # the kth epoch. For now, must be in vector form even if there
    # is only one nonlinear basis function.
    f::BFVT

    # User provided jacobian of f w.r.t. x where h[i,j] ∀ i ∈ [1, k], j ∈ [1, n]
    # computes the partial derivative of f[i] w.r.t. x[j] and takes the state
    # and independant variabels as arguments(i.e. h[i,j](x, tₖ)). For now
    # must be in matrix form even if there is only one nonlinear basis function.
    h::BFJVT

    # Jacobian matrix
    H::Matrix{Float64}

    # Estimate vector
    xhat::Vector{Float64}

    # Measurements
    y::Vector{Float64}

    # Measurement Weights
    W::Symmetric{Float64, Matrix{Float64}}

    # Independant variables
    t::Vector{Float64}

    # Tolerance
    ϵ::Float64

    # Max iterations
    iMax::Int
end

mutable struct ADNonLinearLeastSquares{BFVT} <: AbstractNonLinearLeastSquares
    # Basis functions where f[i] ∀ i ∈ [0, n] should be a function
    # which takes the state and independant
    # variable as arguments (i.e f[i](x,tₖ) where tₖ is the time at
    # the kth epoch.
    f::BFVT

    # Jacobian matrix
    H::Matrix{Float64}

    # Estimate vector
    xhat::Vector{Float64}

    # Measurements
    y::Vector{Float64}

    # Measurement Weights
    W::Symmetric{Float64, Matrix{Float64}}

    # Independant variables
    t::Vector{Float64}

    # Tolerance
    ϵ::Float64

    # Max iterations
    iMax::Int
end

function NonLinearLeastSquares(f::AbstractVector, h::AbstractMatrix, xhat::AbstractVector;
                               ϵ = 1e-6, iMax = 100)
    k = length(f)
    n = length(xhat)
    kh, nh = size(h)
    if k != kh
        throw(ArgumentError("Provided jacobian does not have the correct number of rows."))
    end
    if n != nh
        throw(ArgumentError("Provided jacobian does not have the correct number of columns."))
    end

    NonLinearLeastSquares(f,h,Matrix{Float64}(undef, (0,0)), xhat, Vector{Float64}(undef, 0),
        Symmetric(Matrix{Float64}(undef, (0,0))), Vector{Float64}(undef, 0), ϵ, iMax)
end

function NonLinearLeastSquares(f::AbstractVector, xhat::AbstractVector; ϵ = 1e-6, iMax = 100)
    ADNonLinearLeastSquares(f, Matrix{Float64}(undef, (0,0)), xhat, Vector{Float64}(undef, 0),
        Symmetric(Matrix{Float64}(undef, (0,0))), Vector{Float64}(undef, 0), ϵ, iMax)
end

function FeedMeasurementBatch!(nlls::AbstractNonLinearLeastSquares, ys::AbstractVector, W::AbstractMatrix, ts::AbstractVector)
    m = length(ys)
    n = length(nlls.xhat)
    if m / length(nlls.f) != length(ts)
        throw(ArgumentError("Measurement vector and indepandant variable vector lengths are not appropriate."))
    end
    r, c = size(W)
    if r != c
        throw(ArgumentError("Weight matrix must be square!"))
    end

    # Set measurements and indepandant variable vectors
    nlls.y = ys
    nlls.W = Symmetric(W)
    nlls.t = ts

    # Size H matrix
    nlls.H = zeros(m, n)

    return nothing
end

function FeedMeasurementBatch!(nlls::AbstractNonLinearLeastSquares, ys::AbstractVector, ts::AbstractVector)
    m = length(ys)
    W = zeros(m,m)
    for i in 1:m; W[i,i] = 1.0; end
    FeedMeasurementBatch!(nlls, ys, W, ts)
end

function ComputeEstimate!(nlls::AbstractNonLinearLeastSquares)

    # Get requirements
    conv = false
    n  = length(nlls.xhat)
    ml = length(nlls.t)
    m  = length(nlls.y)
    p  = length(nlls.f)
    iter = 0

    # Allocate memory
    Jco  = Inf
    Δyc  = zeros(m)
    Δx   = zeros(n)
    fxc  = zeros(m)
    tv1  = zeros(m)
    tv2  = zeros(n)
    tm1  = zeros(m,n)
    tm2  = zeros(n,n)
    while conv == false && iter < nlls.iMax
        # Increment iteration counter
        iter += 1

        # Fill H
        fillH!(nlls)

        # Fill expected measurements
        @inbounds for i in 1:ml
            for j in 1:p
                fxc[(i-1)*p + j] = nlls.f[j](nlls.xhat, nlls.t[i])
            end
        end

        # Compute residual and local cost function
        Δyc .= nlls.y .- fxc
        mul!(tv1, nlls.W, Δyc)
        Jc = dot(Δyc, tv1)

        # Compute change in local cost function and print
        δJc = abs(Jc - Jco)
        Jco = Jc
        println("Iteration #: " * string(iter) *
                " δJ: " * string(δJc))

        # Check for convergence
        if δJc < nlls.ϵ / norm(nlls.W)
            conv = true
        else
            # Compute correction
            mul!(tv2, transpose(nlls.H), tv1)
            mul!(tm1, nlls.W, nlls.H)
            mul!(tm2, transpose(nlls.H), tm1)
            fm = factorize(tm2)
            ldiv!(Δx, fm, tv2)

            # Update
            nlls.xhat .+= Δx
        end
    end

    return conv
end

function fillH!(nlls::NonLinearLeastSquares)
    n  = length(nlls.xhat)
    ml = length(nlls.t)
    p  = length(nlls.f)

    # Fill
    @inbounds for i in 1:ml
        for j in 1:p
            for k in 1:n
                nlls.H[(i - 1)*p + j,k] = nlls.h[j,k](nlls.xhat, nlls.t[i])
            end
        end
    end

    return nothing
end

function fillH!(nlls::ADNonLinearLeastSquares)
    n  = length(nlls.xhat)
    ml = length(nlls.t)
    p  = length(nlls.f)

    # Fill
    @inbounds for i in 1:ml
        for j in 1:p
            grad = gradient(x -> nlls.f[j](x, nlls.t[i]), nlls.xhat)[1]
            for k in 1:n
                nlls.H[(i - 1)*p + j,k] = grad[k]
            end
        end
    end

    return nothing
end

function GetMeasurementEstimate(nlls::AbstractNonLinearLeastSquares, t::AbstractFloat)
    return [nlls.f[i](nlls.xhat, t) for i in 1:length(nlls.f)]
end

function GetEstimate(nlls::AbstractNonLinearLeastSquares)
    return nlls.xhat
end


module BernoulliMap

    export Orbit, find_orbits, find_orbits_rand,
        bernoulli_map, orbit_length_minimum, orbit_minimum, orbit_length
    export hofbauer, invariant_measure

    include("hofbauer.jl")
    include("utils.jl")
    include("orbits.jl")
    include("bernoulli_map.jl")
end


@testset "Spin spherical harmonic mode indices (s=$s)" for s in -2:2
    lmax = 10
    N = lmax + 1
    M = 2 * N - 1

    mode_seen = zeros(Int, N, M)
    for l in abs(s):lmax, m in (-l):l
        mode_seen[spinsph_mode(s, l, m)] += 1
    end
    @test mode_seen == spinsphones(Bool, N, M, s)

    # Test some concrete modes
    @test spinsph_mode(0, 0, 0) == CartesianIndex(1, 1)
    @test spinsph_mode(0, 1, -1) == CartesianIndex(1, 2)
    @test spinsph_mode(0, 1, 1) == CartesianIndex(1, 3)

    @test spinsph_mode(1, 1, 0) == CartesianIndex(1, 1)
    @test spinsph_mode(1, 1, -1) == CartesianIndex(1, 2)
    @test spinsph_mode(1, 1, 1) == CartesianIndex(1, 3)

    @test spinsph_mode(2, 2, 0) == CartesianIndex(1, 1)
    @test spinsph_mode(2, 2, -1) == CartesianIndex(1, 2)
    @test spinsph_mode(2, 2, 1) == CartesianIndex(1, 3)
end

@testset "Spin spherical harmonics: simple modes ($T)" for T in
                                                           [Complex{Float64}]
    lmax = 100

    N = lmax + 1
    Θ, Φ = sph_points(N)
    @test length(Θ) == N
    M = length(Φ)

    # Test s=1, l=1, m=0
    s, l, m = 1, 1, 0
    F = T[Y_s1_1_0(θ, ϕ) for θ in Θ, ϕ in Φ]
    @test size(F) == (N, M)
    C = spinsph_transform(F, s)
    @test C ≈ unit(spinsph_mode(s, l, m), size(C))
    F′ = spinsph_evaluate(C, s)
    @test F′ ≈ F
end

@testset "Spin spherical harmonics: simple modes (s=$s, $T)" for s in -2:2,
                                                                 T in [Float64,
                                                                  Complex{Float64}]

    T <: Real && s ≠ 0 && continue

    lmax = 10

    N = lmax + 1
    Θ, Φ = sph_points(N)
    @test length(Θ) == N
    M = length(Φ)

    lmax_test = 4
    for l in abs(s):lmax_test, m in (-l):l
        F = T[sYlm(T, s, l, m, θ, ϕ) for θ in Θ, ϕ in Φ]
        C = spinsph_transform(F, s)
        @test C ≈ unit(spinsph_mode(s, l, m), size(C))
        F′ = spinsph_evaluate(C, s)
        @test F′ ≈ F
    end
end

@testset "Spin spherical harmonics: linearity (s=$s, $T)" for s in -2:2,
                                                              T in [Float64,
                                                               Complex{Float64}]

    T <: Real && s ≠ 0 && continue

    lmax = 100

    N = lmax + 1
    Θ, Φ = sph_points(N)
    @test length(Θ) == N
    M = length(Φ)

    F = zeros(T, N, M)
    C = spinsph_transform(F, s)
    @test iszero(C)

    F1 = randn(T, N, M)
    F2 = randn(T, N, M)
    C1 = spinsph_transform(F1, s)
    C2 = spinsph_transform(F2, s)
    F = F1 + F2
    C = spinsph_transform(F, s)
    @test C ≈ C1 + C2

    F = randn(T, N, M)
    # Why doesn't this work with a complex scale factor? Is the
    # transform not linear?
    # α = randn(T)
    α = randn()
    C = spinsph_transform(F, s)
    Cα = spinsph_transform(α * F, s)
    @test Cα ≈ α * C
end

@testset "Spin spherical harmonics: duality (s=$s, $T)" for s in -2:2,
                                                            T in [Float64,
                                                             Complex{Float64}]

    T <: Real && s ≠ 0 && continue

    lmax = 100

    N = lmax + 1
    Θ, Φ = sph_points(N)
    @test length(Θ) == N
    M = length(Φ)

    F = randn(T, N, M)
    C = spinsph_transform(F, s)
    @test size(C) == size(F)
    @test eltype(C) == eltype(F)
    F′ = spinsph_evaluate(C, s)
    @test F′ ≈ F

    C = randn(T, N, M)
    F = spinsph_evaluate(C, s)
    @test size(F) == size(C)
    @test eltype(F) == eltype(C)
    C′ = spinsph_transform(F, s)
    @test C′ ≈ C
end

@testset "Spin spherical harmonics: orthonormality (s=$s, $T)" for s in -2:2,
                                                                   T in
                                                                   [Float64,
                                                                    Complex{Float64}]

    T <: Real && s ≠ 0 && continue

    lmax = 100
    atol = 4 / lmax^2

    N = lmax + 1
    Θ, Φ = sph_points(N)
    @test length(Θ) == N
    M = length(Φ)

    lmax_test = 4
    for l in abs(s):lmax_test, m in (-l):l
        C = zeros(T, N, M)
        C[spinsph_mode(s, l, m)] = 1
        F = spinsph_evaluate(C, s)
        for l′ in abs(s):lmax_test, m′ in (-l′):l′
            C′ = zeros(T, N, M)
            C′[spinsph_mode(s, l′, m′)] = 1
            F′ = spinsph_evaluate(C′, s)
            int = integrate(F, F′)
            δ = l == l′ && m == m′
            @test isapprox(int, δ; atol=atol)
        end
    end
end

@testset "Spin spherical harmonics: eth and ethbar (s=$s, $T)" for s in -2:2,
                                                                   T in
                                                                   [Complex{Float64}]

    RT = typeof(Real(one(T)))
    lmax = 10

    N = lmax + 1
    Θ, Φ = sph_points(N)
    @test length(Θ) == N
    M = length(Φ)

    lmax_test = 4

    for l in max(abs(s), abs(s + 1)):lmax_test, m in (-l):l
        F = T[ðsYlm(s, l, m, θ, ϕ) for θ in Θ, ϕ in Φ]
        C = spinsph_transform(F, s + 1)
        α = ifelse(m ≥ -s, 1, -1) * sqrt((l - s) * (l + s + 1))
        @test C ≈ unit(α, spinsph_mode(s + 1, l, m), size(C))
        F′ = spinsph_evaluate(C, s + 1)
        @test F′ ≈ F
    end

    for l in max(abs(s), abs(s - 1)):lmax_test, m in (-l):l
        F = T[ð̄sYlm(s, l, m, θ, ϕ) for θ in Θ, ϕ in Φ]
        C = spinsph_transform(F, s - 1)
        α = -ifelse(m > -s, 1, -1) * sqrt((l + s) * (l - s + 1))
        @test C ≈ unit(α, spinsph_mode(s - 1, l, m), size(C))
        F′ = spinsph_evaluate(C, s - 1)
        @test F′ ≈ F
    end

    for l in abs(s):lmax_test, m in (-l):l
        C = zeros(T, N, M)
        C[spinsph_mode(s, l, m)] = ifelse(m ≥ -s, 1, -1)

        ðC = spinsph_eth(C, s)
        ðF = spinsph_evaluate(ðC, s + 1)
        @test ðF .+ 1 ≈ T[ðsYlm(s, l, m, θ, ϕ) for θ in Θ, ϕ in Φ] .+ 1
    end

    for l in abs(s):lmax_test, m in (-l):l
        C = zeros(T, N, M)
        C[spinsph_mode(s, l, m)] = ifelse(m > -s, 1, -1)

        ð̄C = spinsph_ethbar(C, s)
        ð̄F = spinsph_evaluate(ð̄C, s - 1)
        @test ð̄F .+ 1 ≈ T[ð̄sYlm(s, l, m, θ, ϕ) for θ in Θ, ϕ in Φ] .+ 1
    end
end

@testset "Spin spherical harmonics: Laplacian ($T)" for T in [Complex{Float64}]
    lmax = 100

    N = lmax + 1
    Θ, Φ = sph_points(N)
    @test length(Θ) == N
    M = length(Φ)

    # Calculate Laplacian directly
    RT = typeof(Real(one(T)))
    F = randn(RT, N, M)
    C = sph_transform(F)
    ΔC = sph_laplace(C)
    ΔF = sph_evaluate(ΔC)

    # Calculate Laplacian via eth and ethbar
    F⁰ = T.(F)
    C⁰ = spinsph_transform(F⁰, 0)
    @test eltype(C⁰) === T
    ðC¹ = spinsph_eth(C⁰, 0)
    @test eltype(ðC¹) === T
    ð̄ðC⁰ = spinsph_ethbar(ðC¹, 1)
    @test eltype(ð̄ðC⁰) === T
    ð̄ðF⁰ = spinsph_evaluate(ð̄ðC⁰, 0)
    @test eltype(ð̄ðF⁰) === T
    @test ð̄ðF⁰ ≈ ΔF

    # Calculate Laplacian via ethbar and eth
    ð̄C⁻¹ = spinsph_ethbar(C⁰, 0)
    @test eltype(ð̄C⁻¹) === T
    ðð̄C⁰ = spinsph_eth(ð̄C⁻¹, -1)
    @test eltype(ðð̄C⁰) === T
    ðð̄F⁰ = spinsph_evaluate(ðð̄C⁰, 0)
    @test eltype(ð̄ðF⁰) === T
    @test ðð̄F⁰ ≈ ΔF
end

@testset "Spin spherical harmonics: grad ($T)" for T in [Float64]
    lmax = 10

    N = lmax + 1
    Θ, Φ = sph_points(N)
    @test length(Θ) == N
    M = length(Φ)

    lmax_test = 4

    for l in 0:lmax_test, m in (-l):l
        F = Complex{T}[sYlm(0, l, m, θ, ϕ) for θ in Θ, ϕ in Φ]
        ðF₀ = Complex{T}[ðsYlm(0, l, m, θ, ϕ) for θ in Θ, ϕ in Φ]
        ð̄F₀ = Complex{T}[ð̄sYlm(0, l, m, θ, ϕ) for θ in Θ, ϕ in Φ]

        C = spinsph_transform(F, 0)
        @test C ≈ unit(1, spinsph_mode(0, l, m), size(C))

        ðC = spinsph_eth(C, 0)
        ð̄C = spinsph_ethbar(C, 0)

        ðF = spinsph_evaluate(ðC, 1)
        ð̄F = spinsph_evaluate(ð̄C, -1)
        # Different sYlm modes need different signs:
        α = ifelse(m ≥ 0, 1, -1)
        ᾱ = ifelse(m > 0, 1, -1)

        if l == 0
            @test norm(ðF) ≤ sqrt(eps(T))
            @test norm(ð̄F) ≤ sqrt(eps(T))
        else
            @test α * ðF ≈ ðF₀
            @test ᾱ * ð̄F ≈ ð̄F₀
        end
    end

    for l in 0:lmax_test, m in (-l):l
        F = T[Ylm(l, m, θ, ϕ) for θ in Θ, ϕ in Φ]
        ðF₀ = SVector{2,T}[c2a(ðYlm(l, m, θ, ϕ)) for θ in Θ, ϕ in Φ]
        ð̄F₀ = SVector{2,T}[c2a(ð̄Ylm(l, m, θ, ϕ)) for θ in Θ, ϕ in Φ]
        ðC₀ = spinsph_transform(ðF₀, 1)
        ð̄C₀ = spinsph_transform(ð̄F₀, -1)
        ðC₀::Array{SVector{2,T},2}
        ð̄C₀::Array{SVector{2,T},2}

        C = spinsph_transform(F, 0)
        C::Array{T,2}
        for l′ in 0:lmax, m′ in (-l′):l′
            # Allow ±m modes to mix
            if (l′, abs(m′)) ≠ (l, abs(m))
                @test norm(C[spinsph_mode(0, l′, m′)]) ≤ sqrt(eps(T))
            end
        end

        ðC = spinsph_eth(C, 0)
        ð̄C = spinsph_ethbar(C, 0)
        ðC::Array{SVector{2,T},2}
        ð̄C::Array{SVector{2,T},2}

        if l == 0
            @test norm(ðC) ≤ sqrt(eps(T))
            @test norm(ð̄C) ≤ sqrt(eps(T))
        else
            @test ðC ≈ ðC₀
            @test ð̄C ≈ ð̄C₀
        end

        ðF = spinsph_evaluate(ðC, 1)
        ð̄F = spinsph_evaluate(ð̄C, -1)
        ðF::Array{SVector{2,T},2}
        ð̄F::Array{SVector{2,T},2}
        if l == 0
            @test norm(ðF) ≤ sqrt(eps(T))
            @test norm(ð̄F) ≤ sqrt(eps(T))
        else
            @test ðF ≈ ðF₀
            @test ð̄F ≈ ð̄F₀
        end
    end
end

@testset "Spin spherical harmonics: div ($T)" for T in [Float64]
    lmax = 10

    N = lmax + 1
    Θ, Φ = sph_points(N)
    @test length(Θ) == N
    M = length(Φ)

    lmax_test = 4

    for l in 0:lmax_test, m in (-l):l
        ðF = Complex{T}[ðsYlm(0, l, m, θ, ϕ) for θ in Θ, ϕ in Φ]
        ð̄F = Complex{T}[ð̄sYlm(0, l, m, θ, ϕ) for θ in Θ, ϕ in Φ]
        ΔF₀ = Complex{T}[-l * (l + 1) * sYlm(0, l, m, θ, ϕ) for θ in Θ, ϕ in Φ]

        ðC = spinsph_transform(ðF, 1)
        ð̄C = spinsph_transform(ð̄F, -1)
        ð̄ðC = spinsph_ethbar(ðC, 1)
        ðð̄C = spinsph_eth(ð̄C, -1)

        ð̄ðF = spinsph_evaluate(ð̄ðC, 0)
        ðð̄F = spinsph_evaluate(ðð̄C, 0)
        # Different sYlm modes need different signs:
        α = ifelse(m ≥ 0, 1, -1)
        ᾱ = ifelse(m > 0, 1, -1)

        if l == 0
            @test norm(ð̄ðF) ≤ sqrt(eps(T))
            @test norm(ðð̄F) ≤ sqrt(eps(T))
        else
            @test α * ð̄ðF ≈ ΔF₀
            @test ᾱ * ðð̄F ≈ ΔF₀
        end
    end

    for l in 0:lmax_test, m in (-l):l
        ðF = SVector{2,T}[c2a(ðYlm(l, m, θ, ϕ)) for θ in Θ, ϕ in Φ]
        ð̄F = SVector{2,T}[c2a(ð̄Ylm(l, m, θ, ϕ)) for θ in Θ, ϕ in Φ]
        ΔF₀ = T[-l * (l + 1) * Ylm(l, m, θ, ϕ) for θ in Θ, ϕ in Φ]
        ΔC₀ = spinsph_transform(ΔF₀, 0)
        ΔC₀::Array{T,2}

        ðC = spinsph_transform(ðF, 1)
        ð̄C = spinsph_transform(ð̄F, -1)
        ðC::Array{SVector{2,T},2}
        ð̄C::Array{SVector{2,T},2}
        ð̄ðC = spinsph_ethbar(ðC, 1)
        ðð̄C = spinsph_eth(ð̄C, -1)
        ð̄ðC::Array{T,2}
        ðð̄C::Array{T,2}

        if l == 0
            @test norm(ð̄ðC) ≤ sqrt(eps(T))
            @test norm(ðð̄C) ≤ sqrt(eps(T))
        else
            @test ð̄ðC ≈ ΔC₀
            @test ðð̄C ≈ ΔC₀
        end

        ð̄ðF = spinsph_evaluate(ð̄ðC, 0)
        ðð̄F = spinsph_evaluate(ðð̄C, 0)
        ð̄ðF::Array{T,2}
        ðð̄F::Array{T,2}

        if l == 0
            @test norm(ð̄ðF) ≤ sqrt(eps(T))
            @test norm(ðð̄F) ≤ sqrt(eps(T))
        else
            @test ð̄ðF ≈ ΔF₀
            @test ðð̄F ≈ ΔF₀
        end
    end
end

@testset "Spin spherical harmonics: Laplacian ($T)" for T in [Float64]
    lmax = 10

    N = lmax + 1
    Θ, Φ = sph_points(N)
    @test length(Θ) == N
    M = length(Φ)

    for iter in 1:10
        F = randn(T, N, M)

        C = spinsph_transform(F, 0)
        ðC = spinsph_eth(C, 0)
        ð̄C = spinsph_ethbar(C, 0)
        ð̄ðC = spinsph_ethbar(ðC, 1)
        ðð̄C = spinsph_eth(ð̄C, -1)

        for l in 0:lmax, m in (-l):l
            if l == 0
                @test norm(ð̄ðC[spinsph_mode(0, l, m)]) ≤ sqrt(eps(T))
                @test norm(ðð̄C[spinsph_mode(0, l, m)]) ≤ sqrt(eps(T))
            else
                @test ð̄ðC[spinsph_mode(0, l, m)] ≈
                      -l * (l + 1) * C[spinsph_mode(0, l, m)]
                @test ðð̄C[spinsph_mode(0, l, m)] ≈
                      -l * (l + 1) * C[spinsph_mode(0, l, m)]
            end
        end

        ð̄ðF = spinsph_evaluate(ð̄ðC, 0)
        ðð̄F = spinsph_evaluate(ðð̄C, 0)

        C₀ = sph_transform(F)
        ΔC₀ = sph_laplace(C₀)
        ΔF₀ = sph_evaluate(ΔC₀)

        @test ð̄ðF ≈ ΔF₀
        @test ðð̄F ≈ ΔF₀
    end
end


using Turing
using Test

@model gdemo() = begin
  s ~ InverseGamma(2,3)
  m ~ Normal(0, sqrt(s))
  1.5 ~ Normal(m, sqrt(s))
  2.0 ~ Normal(m, sqrt(s))
  return s, m
end

N = 500
s1 = Gibbs(N, HMC(10, 0.1, 5, :s, :m))
s2 = Gibbs(N, PG(10, 10, :s, :m))
s3 = Gibbs(N, PG(10, 2, :s), HMC(1, 0.4, 8, :m))
s4 = Gibbs(N, PG(10, 3, :s), HMC(2, 0.4, 8, :m); thin=false)
s5 = Gibbs(N, CSMC(10, 2, :s), HMC(1, 0.4, 8, :m))

c1 = sample(gdemo(), s1)
c2 = sample(gdemo(), s2)
c3 = sample(gdemo(), s3)
c4 = sample(gdemo(), s4)
c5 = sample(gdemo(), s5)

# Very loose bound, only for testing constructor.
for c in [c1, c2, c3 ,c4, c5]
  check_numerical(c, [:s, :m], [49/24, 7/6], eps=1.0)
end

@test length(c4[:s]) == N * (3 + 2)

# Test gid of each samplers
g = Turing.Sampler(s3)

@test g.info[:samplers][1].alg.gid == 1
@test g.info[:samplers][2].alg.gid == 2


using LinearAlgebra, FastSolversForWeightedTV, CUDA, Flux, Test, BenchmarkTools
CUDA.allowscalar(false)
include("./test_utils.jl")

# flag_gpu = true
flag_gpu = false

# Random data
T = Complex{Float64}
n = (256, 256)
y = randn(T, n..., 2); flag_gpu && (y = y |> gpu)
step = real(T)(1e-5)
rtol = real(T)(1e1*step)

# Norm (2-D)
for g = [mixed_norm(2,2,2; T=T), mixed_norm(2,2,1; T=T), mixed_norm(2,2,Inf; T=T)]

    # Proxy
    λ = real(T)(0.5)*norm(y)^2/g(y)
    x = proxy(y, λ, g)

    # Gradient test
    fun = proxy_objfun(λ, g)
    test_grad(fun, y; step=step, rtol=rtol)

    # Projection test
    ε = real(T)(0.1)*g(y)
    x = project(y, ε, g)
    @test g(x) ≈ ε rtol=rtol

    ## Gradient test
    fun = proj_objfun(ε, g)
    test_grad(fun, y; step=step, rtol=rtol)

end

# Random data
n = (256, 256, 256)
y = randn(T, n..., 3); flag_gpu && (y = y |> gpu)

# Norm (3-D)
for g = [mixed_norm(3,2,2; T=T), mixed_norm(3,2,1; T=T), mixed_norm(3,2,Inf; T=T)]

    # Proxy
    λ = real(T)(0.5)*norm(y)^2/g(y)
    x = proxy(y, λ, g)

    # Gradient test
    fun = proxy_objfun(λ, g)
    test_grad(fun, y; step=step, rtol=rtol)

    # Projection test
    ε = real(T)(0.1)*g(y)
    x = project(y, ε, g)
    @test g(x) ≈ ε rtol=rtol

    ## Gradient test
    fun = proj_objfun(ε, g)
    test_grad(fun, y; step=step, rtol=rtol)

end


"""Will load a collection of 90 problems from a passed filename.
Files must be in the
http://people.brunel.ac.uk/~mastjjb/jeb/orlib/mdmkpinfo.html format."""
function parse_file(filename::String, dataset_num::Int, uid=0)
    f = open(filename)

    problems::Vector{MDMKP_Prob} = []

    #the very first item in the array is the amount of problems found in the
    #file.
    amount_of_problems = next_line(f)[1]

	instance_num = 0
    #so now for every problem:
    for problem in 1:amount_of_problems
		instance_num += 1
        n, m = next_line(f)
        lower_than_values::Vector{Vector{Int}} = []
        for i in 1:m
            push!(lower_than_values, next_line(f))
        end
        lower_than_constraints::Vector{Int} =  next_line(f)
        greater_than_values::Vector{Vector{Int}} = []
        for i in 1:m
            push!(greater_than_values, next_line(f))
        end
        greater_than_constraints = next_line(f)
        cost_coefficient_values::Vector{Vector{Int}} = []
        for i in 1:6
            push!(cost_coefficient_values, next_line(f))
        end

        upper_bounds::Vector{Tuple{Vector{Int},Int}} = []
        lower_bounds::Vector{Tuple{Vector{Int},Int}} = []

        for i in 1:m
            push!(lower_bounds, (greater_than_values[i], greater_than_constraints[i]))
            push!(upper_bounds, (lower_than_values[i], lower_than_constraints[i]))
        end

        q = [1, div(m, 2), m, 1, div(m, 2), m]
        for i in 1:6
			uid += 1
			id = Problem_ID(
				uid,
				dataset_num,
				instance_num,
				i,
				get_tightness(instance_num),
				length(cost_coefficient_values[i]),
				length(lower_bounds[1:q[i]]),
				length(upper_bounds),
				is_mixed(i),
				false,
                '?'
			)
            push!(problems, MDMKP_Prob(
                cost_coefficient_values[i],
                upper_bounds,
                lower_bounds[1:q[i]],
				id
            ))
        end
    end

    (problems, uid)
end

function next_line(file::IOStream)
    return parse_line(readline(file))
end

function parse_line(line)
    return map(parse_int, split(line))
end

function parse_int(x)
    return parse(Int, x)
end

function get_tightness(instance)
	[.25, .5, .75][Int(ceil(instance/5))]
end

function is_mixed(case)
	case > 3
end

function load_folder(
		folder_path="MDMKP/benchmark_problems",
		filename="mdmkp_ct{ds}.txt", datasets=1:9)::Vector{MDMKP_Prob}
	collection = []
	id = 0
	for ds in datasets
		fn = replace(filename, "{ds}"=>"$ds")
		problems, id = parse_file(joinpath(folder_path, fn), ds, id)
		push!(collection, problems)
	end
	vcat(collection...)
end


using OffsetArrays
using .SetProb

function qinit(parm::CParam, meqn::Int, mbc::Int, mx::Int, xlower::Float64, dx::Float64, q::OffsetArray{Float64}, maux::Int, aux::OffsetArray{Float64})

    # Set initial conditions for the q array.
    # This default version prints an error message since it should
    # not be used directly.  Copy this to an application directory and
    # loop over all grid cells to set values of q(1:meqn, 1:mx).

    # integer, intent(in) :: meqn,mbc,mx,maux
    # real(kind=8), intent(in) :: xlower,dx
    # real(kind=8), intent(in) :: aux(maux,1-mbc:mx+mbc)
    # real(kind=8), intent(inout) :: q(meqn,1-mbc:mx+mbc)

    # integer :: i
    # real(kind=8) :: beta, xcell
    # common /cqinit/ beta

    for i = 1 : mx
        xcell = xlower + (i-0.5)*dx
        q[1,i] = exp(-parm.beta * (xcell-0.3)^2)
        q[2,i] = 0.0
        #println("q[1,$i] = $(q[1,i])")
        #println("q[2,$i] = $(q[2,i])")
    end

end # qinit


include(joinpath("joint_degree_seq.jl")) # the joint degree distribution is retrieved from the joint degree sequence

"""
joint_degree_dist(N::UnipartiteNetwork)
    N: Unipartite directed simple network
Returns the proportion of species having a given number of preys and predators
"""
function joint_degree_dist(N::UnipartiteNetwork)
    S = richness(N) # number of species in N
    k = joint_degree_seq(N) # joint degree sequence of N

    # count the number of species having kin predators and kout preys
    k_count = zeros(Int64, S+1, S+1)
    for kin in 0:S
        for kout in 0:S
            k_count[kin+1, kout+1] = sum((k.kin .== kin) .& (k.kout .== kout))
        end
    end

    k_prop = k_count ./ S # proportion of species having kin predators and kout preys
    return k_prop
end


#=
Filename: oligopoly.jl
Author: Victoria Gregory

This is an example for the lecture dyn_stack.rst from the QuantEcon
series of lectures by Tom Sargent and John Stachurski.

We deal with a large monopolistic firm who faces costs:
C_t = e Q_t + .5 g Q_t^2 + .5 c (Q_{t+1} - Q_t)^2

where the fringe firms face:
sigma_t = d q_t + .5 h q_t^2 + .5 c (q_{t+1} - q_t)^2

Additionally, there is a linear inverse demand curve of the form:
p_t = A_0 - A_1 (Q_t + \bar{q_t}) + \eta_t,

where:
    .. math
    \eta_{t+1} = \rho \eta_t + C_{\varepsilon} \varepsilon_{t+1};
    \varepsilon_{t+1} \sim N(0, 1)

For more details, see the lecture.
=#

using QuantEcon
using Roots

"""
This type contains all the parameters and
matrices that describe the oligopoly problem.
"""
struct Oligopoly{TF<:AbstractFloat}
    a0 :: TF
    a1 :: TF
    rho :: TF
    c_eps :: TF
    c :: TF
    d :: TF
    e :: TF
    g :: TF
    h :: TF
    beta :: TF
    A :: Matrix{TF}
    B :: Array{TF}
    Q :: TF
    R :: Matrix{TF}
    Rf :: Matrix{TF}
end

"""
Constructor for building all the elements of the oligopoly
problem.

### Arguments:
* `;a0::AbstractFloat(100)`: Intercept of demand curve
* `;a1::AbstractFloat(1)`: Slope of demand curve
* `;rho::AbstractFloat(0.8)`: Autocorrelation of demand disturbances
* `;c_eps::AbstractFloat(0.2)`: Variance of demand disturbances
* `;c::AbstractFloat(1)`: Cost parameter
* `;d::AbstractFloat(20)`: Cost parameter
* `;e::AbstractFloat(20)`: Cost parameter
* `;g::AbstractFloat(0.2)`: Cost parameter
* `;h::AbstractFloat(0.2)`: Cost parameter
* `;beta::AbstractFloat(0.95)`: Discount rate
"""
function Oligopoly(;a0::AbstractFloat=100.0,
                    a1::AbstractFloat=1.0,
                    rho::AbstractFloat=0.8,
                    c_eps::AbstractFloat=0.2,
                    c::AbstractFloat=1.0,
                    d::AbstractFloat=20.0,
                    e::AbstractFloat=20.0,
                    g::AbstractFloat=0.2,
                    h::AbstractFloat=0.2,
                    beta::AbstractFloat=0.95)

    # Left-hand side of (37)
    Alhs = eye(5)
    Alhs[5, :] = [a0-d 1 -a1 -a1-h c]
    Alhsinv = inv(Alhs)

    # Right-hand side of (37)
    Brhs = [0; 0; 1; 0; 0]
    Arhs = eye(5)
    Arhs[2, 2] = rho
    Arhs[4, 5] = 1
    Arhs[5, 5] = c / beta

    R = [0 0 (a0-e)/2 0 0;
         0 0 1/2 0 0;
         (a0-e)/2 1/2 -a1-0.5*g -a1/2 0;
         0 0 -a1/2 0 0;
         0 0 0 0 0]

    Rf = [0 0 0 0 0 (a0-d)/2;
          0 0 0 0 0 1/2;
          0 0 0 0 0 -a1/2;
          0 0 0 0 0 -a1/2;
          0 0 0 0 0 0;
          (a0-d)/2 1/2 -a1/2 -a1/2 0 -h/2]

    Q = c/2

    A = Alhsinv*Arhs
    B = Alhsinv*Brhs

    return Oligopoly(a0, a1, rho, c_eps, c, d,
                     e, g, h, beta, A, B, Q, R, Rf)

end

"""
Find the value function of the optimal linear regulator problem.
This is steps 2 and 3 in the lecture notes.

### Arguments:
* `olig::Oligopoly`: The oligopoly problem we would like
   to find the value function for

### Returns:
* (P, F, d)::Matrix{TF} where TF<:AbstractFloat : Matrices that describe the value
   function of the optimal linear regulator problem
"""
function find_PFd(olig::Oligopoly)
    lq = LQ(olig.Q, -olig.R, olig.A, olig.B; bet=olig.beta)
    P, F, d = stationary_values(lq)

    Af = vcat(hcat(olig.A-olig.B*F, [0; 0; 0; 0; 0]), [0 0 0 0 0 1])
    Bf = [0; 0; 0; 0; 0; 1]

    lqf = LQ(olig.Q, -olig.Rf, Af, Bf; bet=olig.beta)
    Pf, Ff, df = stationary_values(lqf)

    return P, F, d, Pf, Ff, df
end

"""
Taking the parameters as given, solve for the optimal decision rules
for the firm.

### Arguments:
* `olig::Oligopoly`: The oligopoly problem we would like
   to find the decision rules for
* `(eta0, Q0, q0)::AbstractFloat`: Initial conditions

### Returns:
* (P, -F, D0, Pf, -Ff)::Matrix{Float64}: Matrices that describe the
   decision rule of the optimal linear regulator problem
"""
function solve_for_opt_policy(olig::Oligopoly, eta0::AbstractFloat=0.0,
                              Q0::AbstractFloat=0.0, q0::AbstractFloat=0.0)

    # Step 1/2: Formulate/solve the optimal linear regulator
    P, F, d, Pf, Ff, df = find_PFd(olig)

    # Step 3: Convert implementation into state variables (find coeffs)
    P22 = P[end, end]
    P21 = P[end, 1:end-1]'
    P22inv = P22^(-1)

    # Step 4: Find optimal x_0 and \mu_{x, 0}
    z0 = [1; eta0; Q0; q0]
    x0 = -P22inv*P21*z0
    D0 = -P22inv*P21

    # Return -F and -Ff because we use u_t = -F y_t
    return P, -F, D0, Pf, -Ff

end

olig = Oligopoly()
P, F, D0, Pf, Ff = solve_for_opt_policy(olig)

# Checking time-inconsistency
# Arbitrary initial z_0
y0 = [1; 0; 25; 46]
# optimal x_0 = i_0
i0 = D0*y0
# iterate one period using the closed-loop system
y1 = (olig.A + olig.B*F)*[y0; i0]
# the last element of y_1 is x_1 = i_1
i1_0 = y1[end, 1]

# compare this to the case when the leader solves a Stackelberg problem
# in period 1. if in period 1 the leader could choose i1 given
# (1, v_1, Q_1, \bar{q}_1)
i1_1 = D0*y1[1:end-1, 1]


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

# helper function to extract raw data
# from uniform scaling objects
raw(a::UniformScaling) = a.λ
raw(a) = a

"""
    NestedVariogram(cs, γs)

A nested variogram model `γ = c₁γ₁ + c₂γ₂ + ⋯ + cₙγₙ` with
coefficients `cs = (c₁, c₂, ..., cₙ)` and variogram models
`γs = (γ₁, γ₂, ..., γₙ)`.
"""
struct NestedVariogram{CS,GS} <: Variogram
  cs::CS
  γs::GS

  function NestedVariogram{CS,GS}(cs, γs) where {CS,GS}
    @assert all(issymmetric, cs) "coefficients must be symmetric"
    new(cs, γs)
  end
end

NestedVariogram(cs, γs) = NestedVariogram{typeof(cs),typeof(γs)}(cs, γs)

(g::NestedVariogram)(h)                  = raw(sum(g.cs .* map(γ -> γ(h), g.γs)))
(g::NestedVariogram)(u::Point, v::Point) = raw(sum(g.cs .* map(γ -> γ(u, v), g.γs)))

sill(g::NestedVariogram)         = raw(sum(g.cs .* map(sill, g.γs)))
nugget(g::NestedVariogram)       = raw(sum(g.cs .* map(nugget, g.γs)))
Base.range(g::NestedVariogram)   = maximum(range(γ) for γ in g.γs)
isstationary(g::NestedVariogram) = all(isstationary, g.γs)
isisotropic(g::NestedVariogram) = all(isisotropic, g.γs)

# algebraic structure
*(c, γ::Variogram)                          = NestedVariogram((c,), (γ,))
*(c, γ::NestedVariogram)                    = NestedVariogram(map(x->c.*x, γ.cs), γ.γs)
+(γ₁::Variogram, γ₂::Variogram)             = NestedVariogram((1, 1), (γ₁, γ₂))
+(γ₁::NestedVariogram, γ₂::Variogram)       = NestedVariogram((γ₁.cs..., 1), (γ₁.γs..., γ₂))
+(γ₁::Variogram, γ₂::NestedVariogram)       = NestedVariogram((1, γ₂.cs...), (γ₁, γ₂.γs...))
+(γ₁::NestedVariogram, γ₂::NestedVariogram) = NestedVariogram((γ₁.cs..., γ₂.cs...), (γ₁.γs..., γ₂.γs...))

"""
    structures(γ)

Return the individual structures of a (possibly nested)
variogram as a tuple. The structures are the total nugget
`cₒ`, and the coefficients (or contributions) `cs` for the
remaining non-trivial structures `γs` after normalization
(i.e. sill=1, nugget=0).
"""
function structures(γ::Variogram)
  cₒ = nugget(γ)
  c  = sill(γ) - nugget(γ)
  T  = typeof(c)
  γ  = @set γ.sill = one(T)
  γ  = @set γ.nugget = zero(T)
  cₒ, (c,), (γ,)
end

function structures(γ::NestedVariogram)
  ks, gs = γ.cs, γ.γs

  # total nugget and contributions
  cₒ = raw(sum(@. ks * nugget(gs)))
  cs = @. raw(ks * (sill(gs) - nugget(gs)))

  # discard nugget effect terms
  inds = findall(g->!(g isa NuggetEffect), gs)
  cs, γs = cs[inds], gs[inds]

  # adjust sill and nugget
  γs = map(γs) do γ
    T = typeof(sill(γ))
    γ = @set γ.sill = one(T)
    γ = @set γ.nugget = zero(T)
  end

  cₒ, cs, γs
end

# -----------
# IO METHODS
# -----------

function Base.show(io::IO, g::NestedVariogram)
  O = IOContext(io, :compact => true)
  coeffs = 1 .* raw.(g.cs)
  models = nameof.(typeof.(g.γs))
  lines  = ["$c × $γ" for (c, γ) in zip(coeffs, models)]
  print(O, join(lines, " + "))
end

function Base.show(io::IO, ::MIME"text/plain", g::NestedVariogram)
  O = IOContext(io, :compact => true)
  coeffs = 1 .* raw.(g.cs)
  models = g.γs
  nstruc = length(models)
  name   = "NestedVariogram{$nstruc}"
  header = isisotropic(g) ? name : name*" (anisotropic)"
  println(O, header)
  println(O, "  structures")
  lines = ["    └─$γ" for γ in models]
  print(O, join(lines, "\n"))
  println(O)
  println(O, "  coefficients")
  lines = ["    └─$c" for c in coeffs]
  print(O, join(lines, "\n"))
end


# -*- encoding: utf-8 -*-
#
# The MIT License (MIT)
#
# Copyright © 2021 Matteo Foglieni and Riccardo Gervasoni
#

@testset "test_OrthoNormalBasis" begin
    pcg = PCG()

    for i in 1:100000
        normal = Normal(random(pcg), random(pcg), random(pcg))
        e1, e2, e3 = create_onb_from_z(normal)

        @test e3 ≈ normal # z axis must be aligned with the normal

        # testing orthogonality
        @test are_close(e1 ⋅ e2, 0.)
        @test are_close(e1 ⋅ e3, 0.)
        @test are_close(e3 ⋅ e2, 0.)

        # testing normality
        @test are_close(squared_norm(e1), 1.)
        @test are_close(squared_norm(e2), 1.)
        @test are_close(squared_norm(e3), 1.)
    end
end


# This is a sample setup for a test case with an explicit QDE/DAE
# function.  Fill in the ...

# First , document the test case here!

export mytest

# create an instance of TestCaseExplicit by filling in the ...:
mytest = let
    tcname = ...::Symbol # name should be same as variable name (except for upper/lower case)

    T = ... # the datatype used, probably Float64
    Tarr = [Matrix, SparseMatrixCSC][1] # the container datatype used
                                        # for the Jacobian and mass
                                        # matrix, probably Matrix
                                        # (default) or
                                        # SparseMatrixCSC.
    dof = ...::Int # degrees of freedom
    dae = ...::Int  # index of DAE, ==0 for ODE
    # stiffness of system, one of the three constants
    stiffness = [nonstiff, mildlystiff, stiff][...]

    ## the problem function
    function fn!(t::T,y::Vector{T},dydt::Vector{T})
        # The ode function dydt = f(t,y) modifying dydt in-place.
        # Note that all indices of dydt need to be written!
        ...
        return nothing
    end
    # initializes storage for y:
    fn!( ; T_::Type=T, dof_=dof) = zeros(T_,dof_)

    # NOTE, using keyword arguments more intuitively like:
    # fn!( ; T::Type=T, dof=dof) = zeros(T,dof)
    # does not work: https://github.com/JuliaLang/julia/issues/9948

    ## Jacobian of fn
    function jac!(t::T,y::Vector{T},dfdy::Tarr)
        # the Jacobian of fn, inplace
        #
        # Set jac=nothing if it is not known.
        #
        # Note that jac! should write all indices of given in
        # jpattern.
        ...
        return nothing
    end
    # Returns a matrix which can hold the Jacobian of selected
    # type, initialized to zero.  Can also be used to make a
    # Matlab-style JPattern matrix with jac!(Bool)
    jac!( ; T_::Type=T, dof_=dof) = zeros(T_,dof_,dof_)  # if the problem is large better return an appropriate sparse matrix

    ## Mass matrix
    function mass!(t::T,y::Vector{T},m::Tarr)
        # mass matrix:  M(t,y) * dydt = f(t,y)
        #
        # Set to nothing if it is I, i.e. mass=nothing
        ...
        return nothing
    end
    # Returns a matrix which can hold the mass matrix.
    # Can also be used to make a Matlab-style MvPattern matrix with mass!(Bool)
    mass!( ; T_::Type=T, dof_=dof) = zeros(T_,dof_,dof_)  # if the problem is large better return an appropriate sparse matrix

    ic = T[...] # vector of initial conditions
    tspan = T[start, stop] # integration interval
    refsol = T[...] # reference solution at tspan[2]

    refsolinds = trues(dof)   # if refsol does not contain all indices
                              # then specify which indices of the
                              # solution are to be compared to refsol

    scd_absinds = Int[]       # set where refsol is very small or
                              # zero. To avoid it dominating the
                              # relative error

    tc = TestCaseExplicit{tcname, T, Tarr}(
                             stiffness,
                             dae,
                             dof,
                             fn!,
                             jac!,
                             mass!,
                             ic,
                             tspan,
                             refsol,
                             refsolinds,
                             scd_absinds)

    # put into the right buckets
    tc_all[tcname] = tc
    tc_stiff[tcname] = tc
    # Note that the return-result of the last statement of the let
    # block needs to be tc!
    tc
end


export grad, jacobian, jvp, j′vp, to_vec
replace_arg(x, xs::Tuple, k::Int) = ntuple(p -> p == k ? x : xs[p], length(xs))

"""
    grad(fdm, f, xs...)

Approximate the gradient of `f` at `xs...` using `fdm`. Assumes that `f(xs...)` is scalar.
"""
function grad end

function grad(fdm, f, x::AbstractArray{T}) where T <: Number
    dx = similar(x)
    tmp = similar(x)
    for k in eachindex(x)
        dx[k] = fdm(zero(T)) do ϵ
            tmp .= x
            tmp[k] += ϵ
            return f(tmp)
        end
    end
    return (dx, )
end

grad(fdm, f, x::Real) = (fdm(f, x), )
grad(fdm, f, x::Tuple) = (grad(fdm, (xs...)->f(xs), x...), )

function grad(fdm, f, d::Dict{K, V}) where {K, V}
    dd = Dict{K, V}()
    for (k, v) in d
        function f′(x)
            tmp = copy(d)
            tmp[k] = x
            return f(tmp)
        end
        dd[k] = grad(fdm, f′, v)[1]
    end
    return (dd, )
end

function grad(fdm, f, x)
    v, back = to_vec(x)
    return (back(grad(fdm, x->f(back(v)), v)), )
end

function grad(fdm, f, xs...)
    return ntuple(length(xs)) do k
        grad(fdm, x->f(replace_arg(x, xs, k)...), xs[k])[1]
    end
end

"""
    jacobian(fdm, f, xs::Union{Real, AbstractArray{<:Real}}; len::Int=length(f(x)))

Approximate the Jacobian of `f` at `x` using `fdm`. `f(x)` must be a length `len` vector. If
`len` is not provided, then `f(x)` is computed once to determine the output size.
"""
function jacobian(fdm, f, x::Union{T, AbstractArray{T}}; len::Int=length(f(x))) where {T <: Number}
    J = Matrix{float(T)}(undef, len, length(x))
    for d in 1:len
        gs = grad(fdm, x->f(x)[d], x)[1]
        for k in 1:length(x)
            J[d, k] = gs[k]
        end
    end
    return (J, )
end

function jacobian(fdm, f, xs...; len::Int=length(f(xs...)))
    return ntuple(length(xs)) do k
        jacobian(fdm, x->f(replace_arg(x, xs, k)...), xs[k]; len=len)[1]
    end
end

"""
    _jvp(fdm, f, x::Vector{<:Number}, ẋ::AbstractVector{<:Number})

Convenience function to compute `jacobian(f, x) * ẋ`.
"""
_jvp(fdm, f, x::Vector{<:Number}, ẋ::AV{<:Number}) = fdm(ε -> f(x .+ ε .* ẋ), zero(eltype(x)))

"""
    _j′vp(fdm, f, ȳ::AbstractVector{<:Number}, x::Vector{<:Number})

Convenience function to compute `transpose(jacobian(f, x)) * ȳ`.
"""
_j′vp(fdm, f, ȳ::AV{<:Number}, x::Vector{<:Number}) = transpose(jacobian(fdm, f, x; len=length(ȳ))[1]) * ȳ

"""
    jvp(fdm, f, x, ẋ)

Compute a Jacobian-vector product with any types of arguments for which [`to_vec`](@ref)
is defined.
"""
function jvp(fdm, f, (x, ẋ)::Tuple{Any, Any})
    x_vec, vec_to_x = to_vec(x)
    _, vec_to_y = to_vec(f(x))
    return vec_to_y(_jvp(fdm, x_vec->to_vec(f(vec_to_x(x_vec)))[1], x_vec, to_vec(ẋ)[1]))
end
function jvp(fdm, f, xẋs::Tuple{Any, Any}...)
    x, ẋ = collect(zip(xẋs...))
    return jvp(fdm, xs->f(xs...)[1], (x, ẋ))
end

"""
    j′vp(fdm, f, ȳ, x...)

Compute an adjoint with any types of arguments for which [`to_vec`](@ref) is defined.
"""
function j′vp(fdm, f, ȳ, x)
    x_vec, vec_to_x = to_vec(x)
    ȳ_vec, _ = to_vec(ȳ)
    return (vec_to_x(_j′vp(fdm, x_vec->to_vec(f(vec_to_x(x_vec)))[1], ȳ_vec, x_vec)), )
end
j′vp(fdm, f, ȳ, xs...) = j′vp(fdm, xs->f(xs...), ȳ, xs)[1]

"""
    to_vec(x) -> Tuple{<:AbstractVector, <:Function}

Transform `x` into a `Vector`, and return a closure which inverts the transformation.
"""
to_vec(x::Number) = ([x], first)

# Vectors
to_vec(x::Vector{<:Number}) = (x, identity)
function to_vec(x::Vector)
    x_vecs_and_backs = map(to_vec, x)
    x_vecs, backs = first.(x_vecs_and_backs), last.(x_vecs_and_backs)
    return vcat(x_vecs...), function(x_vec)
        sz = cumsum([map(length, x_vecs)...])
        return [backs[n](x_vec[sz[n]-length(x_vecs[n])+1:sz[n]]) for n in eachindex(x)]
    end
end

# Arrays
to_vec(x::Array{<:Number}) = vec(x), x_vec->reshape(x_vec, size(x))
function to_vec(x::Array)
    x_vec, back = to_vec(reshape(x, :))
    return x_vec, x_vec->reshape(back(x_vec), size(x))
end

# AbstractArrays
function to_vec(x::T) where {T<:LinearAlgebra.AbstractTriangular}
    x_vec, back = to_vec(Matrix(x))
    return x_vec, x_vec->T(reshape(back(x_vec), size(x)))
end
to_vec(x::Symmetric) = vec(Matrix(x)), x_vec->Symmetric(reshape(x_vec, size(x)))
to_vec(X::Diagonal) = vec(Matrix(X)), x_vec->Diagonal(reshape(x_vec, size(X)...))

function to_vec(X::Transpose)
    return vec(Matrix(X)), x_vec->Transpose(permutedims(reshape(x_vec, size(X))))
end
function to_vec(X::Adjoint)
    return vec(Matrix(X)), x_vec->Adjoint(conj!(permutedims(reshape(x_vec, size(X)))))
end

# Non-array data structures

function to_vec(x::Tuple)
    x_vecs, x_backs = zip(map(to_vec, x)...)
    sz = cumsum([map(length, x_vecs)...])
    return vcat(x_vecs...), function(v)
        return ntuple(n->x_backs[n](v[sz[n]-length(x_vecs[n])+1:sz[n]]), length(x))
    end
end

# Convert to a vector-of-vectors to make use of existing functionality.
function to_vec(d::Dict)
    d_vec_vec = [val for val in values(d)]
    d_vec, back = to_vec(d_vec_vec)
    return d_vec, function(v)
        v_vec_vec = back(v)
        return Dict([(key, v_vec_vec[n]) for (n, key) in enumerate(keys(d))])
    end
end


using SimilaritySearch
import StatsBase: fit
using Random
export dnet

"""
A `k`-net is a set of points `M` such that each object in `X` can be:
- It is in `M`
- It is in the knn set of an object in `M` (defined with the distance function `dist`)

The size of `M` is determined by ``\\leftceil|X|/k\\rightceil``

The dnet function uses the `callback` function as an output mechanism. This function is called on each center as `callback(centerId, res)` where
res is a `KnnResult` object (from SimilaritySearch.jl).

"""
function dnet(callback::Function, dist::Function, X::Vector{T}, k) where {T}
    N = length(X)
    metadist = (a::Int, b::Int) -> dist(X[a], X[b])
    I = fit(Sequential, shuffle!(collect(1:N)))
    res = KnnResult(k)
    i = 0

    while length(I.db) > 0
        i += 1
        empty!(res)
        c = pop!(I.db)
        search(I, metadist, c, res)
        callback(c, res)
        @info "computing dnet point $i, dmax: $(covrad(res))"

        j = 0
        for p in res
            I.db[p.objID] = I.db[end-j]
            j += 1
        end

        for p in res
            pop!(I.db)
        end
    end
end


# simple iid bootstrap
function bootstrap(data)
    n = size(data,1)
    resampled = similar(data)
    for i = 1:n
        j = rand(1:n)
        resampled[i,:] = data[j,:]
    end
    return resampled
end


using CSV, Plots
using DiseaseOutbreak

# Hubei and Shandong
hubeiFile = joinpath(@__DIR__, "..", "data",
     "covid19_china","time_series_covid19_confirmed_hubei.csv")
shandongFile = joinpath(@__DIR__, "..", "data",
          "covid19_china","time_series_covid19_confirmed_shandong.csv")

hubei = CSV.read(hubeiFile)
shandong = CSV.read(shandongFile)
hubeiPop = 57.0e6
shandongPop = 94.2e6

hubeiC = convert(Vector{Float64},hubei.ConfirmedCases[1:22])
hubeiFit = fitCaseModel(22,hubeiC,hubeiPop,
                        6.2,8.0,[0.0425,0.0645,2.65])
summary(hubeiFit)
plotFit(hubeiFit,40)
# savefig("hubei.pdf")
pyplotFit(hubeiFit, 40, log10="semilog", grid=false, fsize=[12, 8])
pyplotFit(hubeiFit, 40, log10="loglog", grid=true, fsize=[12, 8])
pyplotFit(hubeiFit, 40, log10="null", grid=true, fsize=[12, 8])

ps = fitCaseModel_ps(22,hubeiC,hubeiPop,
                        6.2,8.0,[0.1,0.1,2.0])
(ps.fit.minimum)

shandongC = convert(Vector{Float64},shandong.ConfirmedCases[1:30])
shandongFit = fitCaseModel(30,shandongC,shandongPop,
                           6.2,8.0,[0.5,0.045,15.0])
shandongFit_ps = fitCaseModel_ps(30,shandongC,shandongPop,
                           6.2,8.0,[0.5,0.045,15.0])
shandongFit_nm = fitCaseModel_nm(30,shandongC,shandongPop,
                           6.2,8.0,[0.5,0.045,15.0])
summary(shandongFit)
shandongFit_ps
shandongFit_nm

plotFit(shandongFit,40)
plot!(estimatedStates(30,shandongPop,shandongC[1],9.66,
                getParams(0.309,0.042,6.2,8.0,SIRX()))[!,:X],
                yaxis=:log,label="paper")

pyplotFit(shandongFit, 40, log10="semilog", grid=true, fsize=[12, 8])
paperRef = estimatedStates(30,shandongPop,shandongC[1],9.66,
                getParams(0.309,0.042,6.2,8.0,SIRX()))[!,:X]

pyplotFit(shandongFit, 30, ref=paperRef, log10="semilog", grid=true,
          fsize=[12, 8])

fitted(shandongFit,40)


"""
File with helper functions
"""

"return U, S, Vt truncated to rank k and the error ||M - M_k||^2"
function truncated_svd(M, k)
    @assert k <= min(size(M)...)
    F = svd(M)
    error = sum(F.S[k+1:end].^2)
    error, SVD(F.U[:,1:k], F.S[1:k], F.Vt[1:k,:])
end

"Returns (indices, elements)"
findmax(v::AbstractArray,k::AbstractUnitRange; opts...) = findpartialsort(v,k; opts..., rev=true)
"Returns (indices, elements)"
findmin(v::AbstractArray,k::AbstractUnitRange; opts...) = findpartialsort(v,k; opts...)

function findpartialsort(v,k; opts...)
    inds = partialsortperm(v,k;opts...)
    length(inds) == 1 ? (inds,v[inds]) : (inds,v[[inds...]])
end

"find n points from data_set closest to the span of U. Returns (indices, elements)"
function find_n_closest_pts(data_set, U, n)
    indices, _ = findmin([distance_squared(x, U) for x in data_set], 1:n)
    indices, data_set[indices]
end

"return square of distance from v to span of U, where columns of U are orthonormal"
distance_squared(v, U) = dot(v,v) - projection_squared(v, U)

"return square of length of projection of v onto the span of U, where columns of U are orthonormal"
projection_squared(v, U) = mapreduce(u->dot(v,u)^2 , +, eachcol(U))

"run method on data_set num_times times. Return errors, indices, and time took for each run"
function run_method(method, data_set, num_times)
    errors = []
    indicess = []
    times = []
    for i in 1:num_times
        t = @elapsed error, indices = method(data_set)
        push!(errors, error)
        push!(indicess, sort(indices))
        push!(times, t)
    end
    return errors, indicess, times
end


using Plots

function level_set(fun; level_center=[0,0], rangeX=[-1,1], rangeY=[-1,1], number_of_levels=50, color_level=true)
    gr()

    centered_RangeX = [level_center[1] + rangeX[1], level_center[1] + rangeX[2]]
    centered_RangeY = [level_center[2] + rangeY[1], level_center[2] + rangeY[2]]

    lengthX = (centered_RangeX[2] - centered_RangeX[1]) / number_of_levels
    x = centered_RangeX[1]:lengthX:centered_RangeX[2]
    lengthY = (centered_RangeY[2] - centered_RangeY[1]) / number_of_levels
    y = centered_RangeY[1]:lengthY:centered_RangeY[2]

    X = repeat(reshape(x, 1, :), length(y), 1)
    Y = repeat(y, 1, length(x))
    Z = map(fun, X, Y)
    p1 = contour(x, y, fun, fill=color_level, c=:starrynight)

    return p1
end

function add_vector(x, y)
    return quiver!([x[1]],[x[2]], quiver=([y[1]],[y[2]]))
end


using IntervalArithmetic
using .DynamicDefinition

"""
Return Branches for a given dynamic, in an iterable
"""
function branches(D::PwMap)
    return [Branch(D.Ts[k], (D.endpoints[k], D.endpoints[k+1]), D.is_full[k] ? (Interval(0),Interval(1)) : (D.Ts[k](Interval(D.endpoints[k]), D.Ts[k](Interval(D.endpoints[k+1])))), D.orientations[k]==1) for k in 1:length(D.Ts)]
end

"""
Return preimages of a certain sequence on all branches of a dynamic
"""
function preimages(seq, D::Dynamic, ϵ = 0.0)
    return [preimages(seq, branch, ϵ) for branch in branches(D)]
end
# TODO: overload for IterateDynamic

"""
Constructs associated "Ulam duals" for a branch b,
i.e., the sequence (k, (T⁻¹(p[k]), T⁻¹(p[k+1]))),
handling endpoints correctly. A callback `f(k, (a, b))` is called on each dual.

This should eventually replace `DualComposedWithDynamic`.
"""
function callback_duals(f, B::Ulam, branch::Branch, preims=nothing)
    if preims === nothing
        preims = preimages(PointSequence(B.p), b)
    end
    n = length(preims.v)

    # duals = Tuple{Int, Tuple{eltype(preims), eltype(preims)}}[]
    # sizehint!(duals, n+2)

    if n==0 # special case: only one preimage
        first_endpoint = preims.increasing ? branch.X[1] : branch.X[2]
        last_endpoint = preims.increasing ? branch.X[2] : branch.X[1]

        f(preims.skip, (first_endpoint, last_endpoint))
        return nothing
    end

    if preims.skip > 0
        first_endpoint = preims.increasing ? branch.X[1] : branch.X[2]
        if first_endpoint != preims.v[1]
            f(preims.skip, (first_endpoint, preims.v[1]))
        end
    end
    for k in 1:n-1
        f(preims.skip+k, (preims.v[k], preims.v[k+1]))
    end
    if preims.skip + n < length(B.p)  # if there are skipped entries at the end of preims
        last_endpoint = preims.increasing ? b.X[2] : b.X[1]
        if last_endpoint != preims.v[n]
            f(preims.skip+n, (preims.v[n], last_endpoint))
        end
    end
    return nothing
end

function callback_duals(f, B::Hat, branch::Branch, preims=nothing)
    if preims === nothing
        preims = preimages(PointSequence(B.p), b)
    end
    @assert preims.skip == 0
    @assert length(preims.v) == length(B.p) # The Hat basis estimation works only for full-branch maps
    for i = 1:length(B) # this skips p[n+1]==1
        x = preims.v[i]
        absT′ = abs(derivative(branch.f, x))
        f(i, (x, absT′))
    end
end

function callback_duals(f, B::Basis, D::Dynamic)
    for (branch, preim) in zip(branches(D), preimages(PointSequence(B.p), D))
        callback_duals(f, B, branch, preim)
    end
end

function assemble2(B::Basis, D::Dynamic, ϵ=0.0; T = Float64)
    # putting types here in hope to improve callback inference
	I::Vector{Int64} = Int64[]
	J::Vector{Int64} = Int64[]
	nzvals::Vector{Interval{T}} = Interval{T}[]
	n::Int64 = length(B)

	# TODO: reasonable size hint?

    function assemble_callback(i, dual_element)
		if !is_dual_element_empty(B, dual_element)
			for (j, x) in ProjectDualElement(B, dual_element)
				push!(I, i)
				push!(J, mod(j,1:n))
				push!(nzvals, x)
			end
		end
    end
    callback_duals(assemble_callback, B, D)

	return sparse(I, J, nzvals, n, n)
end


# script that wants to numerically test chatteris diaconis (misspelled with 99% prob)
# for the estimates of beta, fitness,ergm (many names..) in the DirBin1 case
using Utilities,AReg,StaticNets,JLD,MLBase,StatsBase,DynNets
using PyCall; pygui(:qt); using PyPlot

Nsample = 100
sizes = [50 100  200] #400 800 1200]
Nsizes = length(sizes)
estParSiz = Array{Array{Float64,2},1}(Nsizes)
dgpParSiz = Array{Array{Float64,1},1}(Nsizes)

for k=1:Nsizes
    ##Define the Dgp model
    N = sizes[k]
    deltaN = 5
    maxDeg = sqrt(N)
    tmpdegsIO = round.([ linspace(deltaN,maxDeg,N);Vector(linspace(deltaN,maxDeg,N))])
    dgpPar,dgpDegs, = estimate(ErgmDirBin1(tmpdegsIO))
    #remove infinities
    dgpPar[.!isfinite(dgpPar)] = -3

    dgpParSiz[k] = dgpPar

    L = maximum(abs.(dgpPar))
    C = L
    errBnd = C*sqrt(log(N)/N )
    prob = 1 - C/(N^2)
    dgpMod = ErgmDirBin1(zeros(dgpDegs),dgpPar,Int.(1:N))
    findmax(dgpDegs)
    #sample the dgp model
    sam = sampl(dgpMod,Nsample )

    est = zeros(2N,Nsample)
    degsIO_Sample = [sumSq(sam,2);sumSq(sam,1)]
    estPar = zeros(2N,Nsample)
    #estimate
    for i=1:Nsample
        println(i)
        estPar[:,i]= estimate(fooErgmDirBin1;degIO =degsIO_Sample[:,i])[1]
    end

    estPar[.!isfinite(estPar)] = -3
    estParSiz[k] = estPar
end

k = 1

diff =[ abs.(estParSiz[k] .- dgpParSiz[k]) for k=1:Nsizes]
plt[:hist]([diff[i][:] for i = 1:Nsizes] ,15,density=true )
 title(" A First Test for Asymptotic Consistency    ")
 legend(["N = $(sizes[k])  Bnd = C(L) $(round(sqrt(log(sizes[k])/sizes[k]),3))" for k=1:Nsizes])
  xlabel("θ_est - θ_dgp  ")
 grid(which = "both" , linewidth=1,alpha = 0.5)
##
Mod = ErgmDirBin1(deg,zeros(N),groupsInds)
groupsNames = unique(groupsInds)
NG = length(groupsNames)
tmpDic = countmap(groupsInds)
#number of nodes per group
NGMemb = sortrows([[i for i in keys(tmpDic)] [  tmpDic[i] for i in keys(tmpDic) ] ])
# agrregate degree of each group
degGroups = zeros(2NG) ; for i =1:N degGroups[groupsInds[i]] += deg[i]; degGroups[NG + groupsInds[i]] += deg[N + i];  end
L = sum(degGroups,1)[1]
L/(N^2-N)
groupsPar = ones(2NG)
firstOrderCond(Mod;degIO = deg,parGroupsIO = groupsPar, groupsInds = groupsInds )
expMatrix(Mod,groupsPar)

targetErrValStaticNets = 0.005
estPar,estIt,estMod = estimate(Mod)

##
sam = sampl(estMod,1000)

meanMatSam = squeeze(mean(sam,3),3)
degsSam = expValStatsFromMat(estMod,meanMatSam)
degGroupsSam = zeros(2NG) ; for i =1:N degGroupsSam[groupsInds[i]] += degsSam[i]; degGroupsSam[NG + groupsInds[i]] += degsSam[ N + i] end
relErrSam = maximum(abs.(( degGroupsSam - degGroups )./degGroups )[degGroups.!=0])

## Test of estimates efficiency on the data

## Repeat on piero's DATA
using MAT
pieroEmidEst = matread("/home/Domenico/Dropbox/Dynamic_Networks/data/emid_data/matlabFiles/data\ Piero/forecastEMID/fitness_timeseries.mat")
pieroEmidData = matread("/home/Domenico/Dropbox/Dynamic_Networks/data/emid_data/matlabFiles/data\ Piero/forecastEMID/eMID_weeklyaggregated_binaryDIRECT_postLTRO.mat")

pieroA_T = Bool.(pieroEmidData["A"])

# load my data
halfPeriod = false
## Load dataj
fold_Path =  "/home/Domenico/Dropbox/Dynamic_Networks/data/emid_data/juliaFiles/"
loadFilePartialName = "Weekly_eMid_Data_from_"
halfPeriod? periodEndStr =  "2012_03_12_to_2015_02_27.jld": periodEndStr =  "2009_06_22_to_2015_02_27.jld"
@load(fold_Path*loadFilePartialName*periodEndStr, AeMidWeekly_T,banksIDs,inactiveBanks, YeMidWeekly_T,weekInd,datesONeMid ,degsIO_T,strIO_T)

matA_T = AeMidWeekly_T# pieroA_T
degsIO_T = [sumSq(matA_T,2);sumSq(matA_T,1)]
T = size(matA_T)[3]
allFitSS =  StaticNets.estimate( StaticNets.SnapSeqNetDirBin1(degsIO_T); identPost = false,identIter= true )
dgpPar = meanSq(allFitSS,2 )
expDgpMat = StaticNets.expMatrix(fooErgmDirBin1,dgpPar)
dgpDegs = [sumSq(expDgpMat,2); sumSq(expDgpMat,1)]
N = length(matA_T[1,:,1])
Nsample = 1000
L = maximum(abs.(dgpPar))
C = L
errBnd = C*sqrt(log(N)/N )
prob = 1 - C/(N^2)
dgpMod = ErgmDirBin1(zeros(dgpDegs),dgpPar,Int.(1:N))
findmax(dgpDegs)
#sample the dgp model
sam = sampl(dgpMod,Nsample )

est = zeros(2N,Nsample)
degsIO_Sample = [sumSq(sam,2);sumSq(sam,1)]
estPar = zeros(2N,Nsample)
estDegs = zeros(2N,Nsample)
#estimate
for i=1:Nsample
    println(i)
    estPar[:,i]= estimate(fooErgmDirBin1;degIO =degsIO_Sample[:,i])[1]
    expEstMat = StaticNets.expMatrix(fooErgmDirBin1,estPar[:,i])
    estDegs[:,i] = [sumSq(expEstMat,2); sumSq(expEstMat,1)]

end

estPar[.!isfinite(estPar)] = -5

diffdat = estDegs .- dgpDegs # abs.(estPar .- dgpPar)
plt[:hist](diffdat[:]  ,100,density=true )
 title("Sample  Static Model, Run MLE and Conpare Degrees")
 grid(which = "both" , linewidth=1,alpha = 0.5)
 xlim(-3, 3)
 xlabel("deg_dgp - deg_MLE")

sampMeanDegs = meanSq(degsIO_Sample,2)
ind1= findmin(sampMeanDegs)[2]
ind2 = findmax(sampMeanDegs)[2]
diffdat = abs.(estPar .- dgpPar)./abs.(dgpPar)
plt[:hist]([diffdat[ind1,:] diffdat[ind2,:]]  ,50,density=true )
 title("DGP fit = mean sing snap estimates  ")
 legend(["Sample Mean Deg = $(round(sampMeanDegs[ind1],2)) " ; "Sample Mean Deg = $(round(sampMeanDegs[ind2],2)) "])
  xlabel("θ_est - θ_dgp  ")
 grid(which = "both" , linewidth=1,alpha = 0.5)

degsIO_Sample[1,:]
degsIO_Sample[ind2,:]

funProb(x::Array{T,1} where T<: Real ) = 1./(1+exp.(-x))
derfunProb(x::Array{T,1} where T<: Real )  =  exp.(-x)./((1+exp.(-x)).^2)
x = Vector(-5:0.001:5)
othPar = -3:1.5:3
for i = 1:length(othPar) plot(x, funProb(x.+othPar[i] )); end
 legTex = ["θ_2 =  $(othPar[i])" for i = 1:length(othPar)]
 title( "1/(1+ exp( - θ - θ_2 ) )")
 legend(legTex)
 xlabel("θ")
 grid()

#


@testset "utils" begin
    @test @inferred(TreeArrays.nfill(2, Val(2))) === (2,2)
    @test @inferred(TreeArrays.nfill(2, Val(3))) === (2,2,2)
end

@testset "Power2" begin
    @test TreeArrays.Power2(0) == 2^0
    @test TreeArrays.Power2(1) == 2^1
    @test TreeArrays.Power2(2) == 2^2

    x = TreeArrays.Power2(3)
    x′ = 2^3
    @test @inferred(2 * x) === 2 * x′
    @test @inferred(x * 2) === x′ * 2
    @test @inferred(x * x) === TreeArrays.Power2(6) == x′ * x′

    @test @inferred(x + 1) === x′ + 1
    @test @inferred(x - 1) === x′ - 1
    @test @inferred(1 + x) === 1 + x′
    @test @inferred(1 - x) === 1 - x′

    @test @inferred(div(20, x)) === div(20, x′)
    @test @inferred(rem(20, x)) === rem(20, x′)
    @test @inferred(divrem(20, x)) === divrem(20, x′)

    @test @inferred(zero(x)) === TreeArrays.Power2(0)
end


function sieve(limit)
    candidates = 2:limit
    out = Int[]
    while length(candidates) > 0
        x = candidates[1]
        push!(out, x)
        candidates = filter(c -> (c .% x) .!= 0, candidates)
    end
    return out
end


""" Convert a presentation math expression to a MML string or a list of MML string
"""
exlist_to_mml(nodes::Vector{CompoExpr}) = foldl(*, map(expr_to_mml, nodes), init="")

expr_to_mml(n::Sub) = mml_element(:msub, _2_exprs_to_mml(n.base, n.sub))
expr_to_mml(n::Sup) = mml_element(:msup, _2_exprs_to_mml(n.base, n.sup))
expr_to_mml(n::SubSup) = mml_element(:msubsup, _3_exprs_to_mml(n.base, n.sub, n.sup))
expr_to_mml(n::Under) = mml_element(:munder, _2_exprs_to_mml(n.base, n.under))
expr_to_mml(n::Over) = mml_element(:mover, _2_exprs_to_mml(n.base, n.over))
expr_to_mml(n::UnderOver) = mml_element(
        :munderover,
       _3_exprs_to_mml(n.base, n.under, n.over))
expr_to_mml(n::Sqrt) = mml_element(:msqrt, expr_to_mml(n.base))
expr_to_mml(n::Frac) = mml_element(:mfrac, _2_exprs_to_mml(n.num, n.den))
expr_to_mml(n::Group) = mml_element(:mrow, exlist_to_mml(n.val))
# we do not really use the `mfenced` element yet. Even firefox does not
# support it as of 2021 (but MathJax does). Per MathML 3.0 specification (2nd
# edition), the following construct is equivalent to `mfenced`.
expr_to_mml(n::Fenced) = mml_element(:mrow,
    mml_delim_pairs(n.ldelim, n.rdelim, exlist_to_mml(n.val)))
# MML Binom is a hack, there is no "built-in" MML element for a Binom
expr_to_mml(n::Binom) = mml_element(
    :mrow,
    mml_delim_pairs(
        '(', ')',
        mml_element_attr(
            :mfrac,
            _2_expr_to_mml(n.n, n.k),
            Dict("linethickness"=>"0")
        )))

function expr_to_mml(n::Environment)
    mat_delim = Dict("pmatrix"=>"()",
                    "bmatrix"=>"[]",
                    "vmatrix"=>"||",
                    "Vmatrix"=>"∥∥")
    if n.tag == "matrix"
        return make_mml_table(n.val)
    elseif n.tag in ("pmatrix", "bmatrix", "vmatrix")
        return mml_element(
            :mrow,
            mml_delim_pairs(mat_delim[n.tag][1], mat_delim[n.tag][2], make_mml_table(n.val)))
    elseif n.tag == "verbatim"
        # We quietly ignore the `verbatim` environment for now.
        # Perhaps we don't even need a `:mrow` tag, just return the content
        #return mml_element(:mrow, exlist_to_mml(n.val))
        return exlist_to_mml(n.val)
    else
        @error("Unknown environment: "*n.tag)
    end
end

function make_mml_table(m::Vector{CompoExpr})
    # t: table, r: row, c: cell
    r, c, t = "", "", ""
    for n in m
        if n isa Ampersand
            # create a mml cell whose content is `c`
            r *= mml_element(:mtd, c)
            c = ""
        elseif n isa Newline
            # create a mml row in `r`
            r *= mml_element(:mtd, c)
            t *= mml_element(:mtr, r)
            r, c = "", ""
        else
            # amalgamate elements in a cell
            c *= expr_to_mml(n)
        end
    end
    # finish the last row (assume there is no ampersand or newline
    # behind the last expression in the input vector)
    r *= mml_element(:mtd, c)
    t *= mml_element(:mtr, r)
    return mml_element(:mtable, t)
end

expr_to_mml(n::ID) = mml_element(:mi, n.val)
expr_to_mml(n::NUM) = mml_element(:mn, n.val)
expr_to_mml(n::OPID) = mml_element(:mo, n.val)
expr_to_mml(n::SPACE) = mml_element_attr(:mspace, "", Dict("width"=>n.val))
expr_to_mml(n::Ampersand) = mml_element(:mo, "&")
# The following should be the correct implementation, but as of 2021/7/1
# Firfox does not support the `linebreak="newline"` attribute.
#   "<mspace linebreak=\"newline\" width=\"2ex\" />"
#
# A dirty hack here. the <math> tag is wrapped by LaTeX2MathML.show().
# But we have no idea if the display style is `inline` or `block` here.
expr_to_mml(n::Newline) = "</math><br/><math>"
expr_to_mml(n) = mml_element(:unknown, repr(n))

mml_element(ele, content) = string("<", ele, ">", content, "</", ele, ">")

function mml_element_attr(ele, content, attr::Dict)
    s = string("<", ele, " ")
    for k in keys(attr)
        s *= k*"=\""*attr[k]*"\" "
    end
    s *= ">"*content*"</"*string(ele)*">"
    return s
end

function mml_element_attr(ele, content, attr::String)
    return string("<", ele, " ", attr, ">", content, "</",  ele, ">")
end

mml_delim_pairs(open, close, content) = (
    mml_element(:mo, open)*content*mml_element(:mo, close))

_2_exprs_to_mml(c1, c2) = expr_to_mml(c1)*expr_to_mml(c2)
_3_exprs_to_mml(c1, c2, c3) = _2_exprs_to_mml(c1, c2)*expr_to_mml(c3)

#############################################################################
## pretty_print_mml() is for debug
function pretty_print_mml(mml::String)
    # convert a utf-8 string to a vector of unicode chars so we can
    # index into it easily
    mml_chars = Char[]
    i = 1
    while i <= ncodeunits(mml)
        push!(mml_chars, mml[i])
        i = nextind(mml, i)
    end
    ele_stack = String[]
    t = 1
    while t <= length(mml_chars)
        #skip whitespaces
        while isspace(mml_chars[t]) t += 1 end
        mml_chars[t] != '<' && error("Unknown token: ", mml_chars[t])
        t = pretty_print_element(mml_chars, t, 0, ele_stack)
    end
end

const tbsz = 3
function pretty_print_element(mml::Vector{Char}, start::Int, indent::Int, ele_stack::Vector{String})
    # mml[start] must be '<'
    i, ele_name = start + 1, ""
    # get starting tag
    while mml[i] != '>'
        ele_name *= mml[i]
        i += 1
    end
    i += 1
    push!(ele_stack, ele_name)
    print(" "^indent*"<", ele_name, '>')
    if ele_name in ("mi", "mn", "mo", "mspace")
        # simple leaf elements
        while mml[i] != '<'
            print(mml[i])
            i += 1
        end
        # Now mml[i] == '<' in the ending tag of <mi>, <mn>, <mo>, <mspace>
        while mml[i] != '>' i += 1 end
        i += 1
        i <= length(mml) && while isspace(mml[i]) i += 1 end
        print("</", ele_name, ">")
        println()
        pop!(ele_stack)
        return i
    else
        # the starting tag name occupies a line
        println()
        # keep finding the next tag and print
        while true
            while isspace(mml[i]) i += 1 end
            if i > length(mml)
                break
            end
            if mml[i] == '<' && mml[i+1] != '/'
                i = pretty_print_element(mml, i, indent + tbsz, ele_stack)
            elseif mml[i+1] == '/'
                i += 2
                en2 = ""
                while mml[i] != '>'
                    en2 *= mml[i]
                    i += 1
                end
                en = pop!(ele_stack)
                # we could have sth like <mspace width=...> </mspace>
                !startswith(en, en2) && begin
                    println(ele_stack)
                    error("tag not match: ", en, " -- ", en2)
                end
                # mml[i] == '>'
                println(" "^indent*"</", en2, '>')
                i += 1
                i > length(mml) && return i
                while isspace(mml[i]) i += 1 end
                mml[i] != '<' && error("unknown tag: ", mml[i])
                return i
            else
                error("unknown tag: ", mml[i])
            end
        end
    end
end

###############################################################################
"""Convert a presentation math expression to a LaTeX string"""
exlist_to_tex(nodes::Vector{CompoExpr}) = foldl(*, map(expr_to_tex, nodes), init="")

# unicode identifiers are kept unchanged
expr_to_tex(ex::ID) = ex.val
expr_to_tex(ex::NUM) = ex.val
expr_to_tex(ex::OPID) = ex.val
expr_to_tex(ex::SPACE) = raw"\,"
expr_to_tex(::Ampersand) = "&"
expr_to_tex(::Newline) = "\\\\"

expr_to_tex(n::Sub) = "{$(expr_to_tex(n.base))}_{$(expr_to_tex(n.sub))}"
expr_to_tex(n::Sup) = "{$(expr_to_tex(n.base))}^{$(expr_to_tex(n.sup))}"
expr_to_tex(n::SubSup) =
    "{$(expr_to_tex(n.base))}_{$(expr_to_tex(n.sub))}^{$(expr_to_tex(n.sup))}"
expr_to_tex(n::Under) = "\\underset{$(n.under)}{$(n.base)}"
expr_to_tex(n::Over) = "\\underset{$(n.under)}{$(n.base)}"
expr_to_tex(n::UnderOver) =
    "\\underset{$(n.under)}{\\overset{$(n.over)}{$(n.base)}}"
expr_to_tex(n::Sqrt) = "\\sqrt{$(n.base)}"
expr_to_tex(n::Frac) = "\\frac{$(n.base)}"
expr_to_tex(n::Group) = "{$(exlist_to_tex(n.val))}"
expr_to_tex(n::Fenced) = "\\left$(n.ldelim)$(exlist_to_tex(n.val))\\right$(n.rdelim)"
expr_to_tex(n::Binom) = "\\binom{$(n.n)}{$(n.k)}"

function expr_to_tex(n::Environment)
    return "\\begin{$(n.tag)}"*exlist_to_tex(n.val)*"\\end{$(n.tag)}"
end


"""
    Interpolate <: Imputor

Performs linear interpolation between the nearest values in an vector.
"""
struct Interpolate <: Imputor end

"""
    impute!(imp::Interpolate, ctx::Context, data::AbstractVector)

Uses linear interpolation between existing elements of a vector to fill in missing data.

WARNING: Missing values at the head or tail of the array cannot be interpolated if there
are no existing values on both sides. As a result, this method does not guarantee
that all missing values will be imputed.
"""
function impute!(imp::Interpolate, ctx::Context, data::AbstractVector{<:Union{T, Missing}}) where T
    i = findfirst(ctx, data) + 1

    while i < length(data)
        if ismissing(ctx, data[i])
            prev_idx = i - 1
            next_idx = findnext(ctx, data, i + 1)

            if next_idx !== nothing
                gap_sz = (next_idx - prev_idx) - 1

                diff = data[next_idx] - data[prev_idx]
                incr = diff / T(gap_sz + 1)
                start_val = data[prev_idx]
                stop_val = data[next_idx]

                values = Real(start_val):Real(incr):Real(stop_val)

                idx_range = prev_idx:(prev_idx + length(values) - 1)
                # println(collect(idx_range))
                # println(values)

                data[idx_range] = values
                i = next_idx
            else
                break
            end
        end
        i += 1
    end

    return data
end


"""
`c, σ2 = get_centers_linear(Nbasis,x)`\n
`x` is the (multidimensional) scheduling signal
"""
function get_centers_linear(Nbasis,x)
    ma = maximum(x,dims=1)
    mi = minimum(x,dims=1)
    n = size(x,2)
    d = ma-mi
    Ni = d./(Nbasis+1)

    σ2 = zeros(Nbasis,n)
    c = zeros(Nbasis,n)
    for i = 1:n
        σ2[:,i] .= d[i]*(0.5/Nbasis)^2
        c[:,i]  = LinRange(mi[i]+0Ni[i],ma[i]-0Ni[i],Nbasis)
    end
    return c, σ2
end

# function get_centers_log(Nbasis)
#     Ni = 1/(Nbasis+1)
#     return (logspace(log10(Ni),log10(1-Ni),Nbasis))[end:-1:1]
# end

comp(x) = (x)

function kernel_matrix(x::AbstractVecOrMat,c,σ2)
    Ψ = Float64[exp(-1/(2σ2[j])*(comp(x)-(c[j]))^2) for x in x, j in eachindex(c)]
    Ψ ./= sum(Ψ,dims=2)
    Ψ
end

function kernel_vector(x::Real,c,σ2)
    Ψ = Float64[exp(-1/(2σ2[j])*(comp(x)-(c[j]))^2) for j in eachindex(c)]
    Ψ ./= sum(Ψ)
    Ψ
end

function kernel_vector(dmp::AbstractDMP, x::Real)
    kernel_vector(x, dmp.c, dmp.σ2)
end

function centraldiff(v::AbstractMatrix)
    dv = diff(v,dims=1)./2
    a1 = [dv[[1],:];dv]
    a2 = [dv;dv[[size(dv,1)],:]]
    a = a1+a2
end

function centraldiff(v::AbstractVector)
    dv = diff(v)./2
    a1 = [dv[1];dv]
    a2 = [dv;dv[end]]
    a = a1+a2
end

_1(y::VecOrMat) = y[1,:][:]
_1(dmp::AbstractDMP) = _1(dmp.y)
_T(dmp::AbstractDMP) = size(dmp.y,1)

math(sl) = identity(sl) #map(s->string("\$",s,"\$") ,sl)

function euler(time_derivative, state0, t, args...; kwargs...)
    T = length(t)
    n = length(state0)
    res = Matrix{Float64}(undef,T,n)
    res[1,:] = state0
    for i in 2:T
        res[i,:] = res[i-1,:] + time_derivative(t[i-1],res[i-1,:])*(t[2]-t[1])
    end
    t,res
end

"""
plot(dmp::DMP, phase::Bool=false; [y0])
"""
@recipe function plotdmp(dmp::AbstractDMP, phase::Bool=false; y0 = _1(dmp))
    n = size(dmp.y,2)
    if n < 2 && phase
        warn("Phase plotting only supported for DMP of dimension greater than 1")
        phase = false
    end
    layout := (n,1)
    tout,yout,ẏout,xout = solve(dmp, y0=y0)[1:4]
    delete!(plotattributes,:y0)
    if phase
        xguide := "Position 1"
        yguide := "Position 2"
    else
        xguide := "Time"
        yguide := "Output"
    end
    for i = 1:n
        @series begin
            title       := "Dynamic Movement Primitive $i"
            subplot     := i
            seriestype  := :path
            label       := ["y_{out}" "ẏ_{out}"] |> math
            if phase
                xguide := "Position 1"
                yguide := "Position 2"
                (yout[:,1], yout[:,2])
            else
                xguide := "Time"
                yguide := "Output"
                (tout, [yout[:,i] ẏout[:,i]])
            end

        end
        @series begin
            title       := "Dynamic Movement Primitive $i"
            subplot     := i
            seriestype  := :path
            linestyle   := :dash
            label       := ["y" "ẏ"] |> math
            phase ? (dmp.y[:,1], dmp.y[:,2]) : (tout, [dmp.y[:,i] dmp.ẏ[:,i]])
        end
    end
    delete!(plotattributes,:phase)
    nothing
end


"""
Algorithm for finding k-cliques (complete subgraphs of order k).

Based on "Arbority and Subgraph Listing Algorithms" by Norshige Chiba and Takao
Nishizeki. The original paper may be found at
https://pdfs.semanticscholar.org/0d19/245a27bc65a87a8014d5b8a66fb514c8ff0b.pdf?_ga=2.251120710.1988770043.1596204199-972144793.1596204199.

Code copied from https://github.com/arbenson/HigherOrderClustering.jl/blob/master/src/cliques.jl
"""

using MatrixNetworks
using LinearAlgebra
using SparseArrays

function _kcliques(A::SparseMatrixCSC{T,Int64}, k::Int64) where T
    A = min.(A, 1)
    A -= Diagonal(A)
    # "global" variables for the recursive algorithm
    n = size(A, 2)
    adjacency_lists = Vector{Vector{Int64}}()
    for v = 1:n
        push!(adjacency_lists, filter(x -> x != v, findnz(A[:, v])[1]))
    end
    C = Int64[]
    labels = k * ones(Int64, n)
    counts = zeros(Int64, n)
    cliques = []

    # Get all neighbors of a node with a given label
    function neighbors(node::Int64, label::Int64)
        index = 1
        for nbr in adjacency_lists[node]
            nbr_label = labels[nbr]
            if nbr_label == label; index += 1; end
            if nbr_label  > label; break; end
        end
        return adjacency_lists[node][1:(index - 1)]
    end

    # Get the number of neighbors of a node with a given label
    function num_neighbors(node::Int64, label::Int64)
        index = 1
        for nbr in adjacency_lists[node]
            nbr_label = labels[nbr]
            if nbr_label == label; index += 1; end
            if nbr_label  > label; break; end
        end
        return index - 1
    end

    # Get the nodes in U in sorted order that need to be processed
    function nodes_to_process(U::Vector{Int64}, label::Int64)
        sorted_nbrs = sort([(num_neighbors(v, label), v) for v in U], rev=true)
        return [x[2] for x in sorted_nbrs if x[1] > 1]
    end

    # Adjust adjacency lists of neighbors of nodes in U to put neighbors with a
    # given label in the front.
    function move_front(U::Vector{Int64}, label::Int64)
        for v in U
            front_nbrs = Int64[]
            back_nbrs  = Int64[]
            for nbr in adjacency_lists[v]
                if labels[nbr] == label; push!(front_nbrs, nbr)
                else                     push!(back_nbrs, nbr)
                end
            end
            num_front = length(front_nbrs)
            total = length(adjacency_lists[v])
            adjacency_lists[v][1:num_front] = front_nbrs
            adjacency_lists[v][(num_front + 1):total] = back_nbrs
        end
    end

    # Adjust adjacency lists of neighbors of nodes in U to put node vi directly
    # after nodes with a given label.
    function move_after(U::Vector{Int64}, vi::Int64, label::Int64)
        for v in U
            front_nbrs = Int64[]
            for nbr in adjacency_lists[v]
                if nbr == vi; continue; end
                if labels[nbr] <= label; push!(front_nbrs, nbr)
                else                     break
                end
            end
            num_front = length(front_nbrs)
            adjacency_lists[v][1:num_front] = front_nbrs
            adjacency_lists[v][num_front + 1] = vi
        end
    end

    # Process all current cliques, which each consist of:
    #   1. the nodes in the global stack C
    #   2. a node in U
    #   3. a neighbor of u with label 2
    function process_cliques(U::Vector{Int64})
        C_cnt = 0
        for v in U
            v_cnt = 0
            for nbr in neighbors(v, 2)
                # Only count (nbr, v) edge once.
                if nbr > v
                    counts[nbr] += 1
                    v_cnt += 1
                    push!(cliques, sort([[nbr, v]; C]))
                end
            end
            counts[v] += v_cnt
            C_cnt += v_cnt
        end
        counts[C] .+= C_cnt
    end

    # Recursive algorithm of Chiba and Nishizeki.
    function rcliques(U::Vector{Int64}, r::Int64)
        # Check base case
        if r == 2
            process_cliques(U)
            return
        end

        for (ind, v) in enumerate(nodes_to_process(U, r))
            # Get neighborhood of v
            Up = neighbors(v, r)
            labels[Up] .= r - 1
            move_front(Up, r - 1)

            # Recurse on neighborhood of node v
            push!(C, v)
            rcliques(Up, r - 1)
            pop!(C)

            # Restore neighborhood
            labels[Up] .= r

            # Eliminate v
            labels[v] = r + 1
            move_after(Up, v, r)
        end
    end

    rcliques(findall(vec(sum(A, dims=2)) .>= (k - 1)), k)
    return counts, cliques
end

"""
`kcliques`
===========
The Chiba-Nishizeki algorithm for finding all k-cliques in a graph. Found on
pp. 216 of the paper. Code from https://github.com/arbenson/HigherOrderClustering.jl

Arguments
---------
    - `A::MatrixNetwork`: The graph in which to search for k-cliques
    - `k::Int64`: The desired clique size

The algorithm takes advantage of the following observations:
    1) a node is in a k-clique iff the subgraph induced by its neighbors
       contains a (k-1)-clique.
    2) if v is in the subgraph "A" induced by the neighbors of u, and w is in
       the subgraph induced by the neighbors of v in "A", then [u,v,w] is a clique.
So, the algorithm recursively finds k-cliques by exploring subgraphs of subgraphs
of subgraphs ... of neighbors, taking note of which vertices remain after k
iterations.
"""
function kcliques(A::MatrixNetwork, k::Int64) where T
    S = sparse(A)
    S = min.(S, 1)
    S -= Diagonal(S)
    clique_counts = zeros(Int64, size(S, 1))
    (d, rt) = corenums(S)
    inds = findall(d .>= (k - 1))
    counts, cliques = _kcliques(S[inds, inds], k)
    clique_counts[inds] = counts
    return clique_counts, cliques
end


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

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ ea8cc0e4-9920-11eb-360c-a1ca0e5fea2a
begin
	using Pluto,PlutoUI,Plots
	using Printf: @sprintf
	plotly()
end

# ╔═╡ 483ce98d-3409-4fed-a360-89291d643ae1
md"""# Bienvenide al Pequeño Graficador de Sucesiones

La idea de esta herramienta es que juegues con los controles para ir explorando las distintas caracteristicas de las sucesiones.

A divertirse :)

tip: los gráficos son interactivos
"""

# ╔═╡ 12f9cc2e-e335-4242-89ac-e3819c944266
md"N: $(@bind N Slider(0:100,default=1,show_value=true)) es la cantidad de puntos a graficar"

# ╔═╡ a1b36581-d9fc-4b41-8c41-1a6d64d95d4a
md"Definir la sucesión $a_{n}$ en la siguiente:"

# ╔═╡ 742547cf-890a-4b6f-8d24-92772b5b3f14
a(n) = 1/n

# ╔═╡ 3deee887-5ec0-4b01-84bb-b1a4edd0cffb
scatter(a.(1:N),zeros(N),hover=["n:$i, an:$(@sprintf("%.5g",a(i)))" for i in 1:N],ylim=(-0.01,0.01),marker=:vline,yaxis=false,legend=false)

# ╔═╡ 52c81bda-00c9-4c68-ab02-8a78dd8e79a0
md"""### Límite de una sucesión
	Sea (aₙ)  n ∈ N una sucesión de números reales.
	Decimos que un número real *L* es límite de la sucesión si para todo ϵ > 0 existe	un natural n₀ = n₀(ϵ) a partir del cual, los valores de la sucesión distan de *L* en menos de ϵ.

	En símbolos

		∀ϵ > 0, ∃ n₀ = n₀(ϵ) ∈ N tal que ∀n > n₀ ⇒ |L − aₙ| < ϵ
	"""

# ╔═╡ 35838740-5278-4e35-bc1f-58e783789101
md"""
L = $(@bind L Slider(-10:1:10,default=1,show_value=true))

ϵ₀ = $(@bind ϵ0 Slider(0:0.01:1,default=1,show_value=true))
E  = $(@bind E Slider(-10:1,default=0,show_value=true))

nᵢ = $(@bind n0 Slider(0:1000,show_value=true))   *n_inicial*
"""

# ╔═╡ 13fce500-346f-4446-a6b5-9c5270c25d6f
begin
	ϵ = ϵ0*10.0^E
	md"ϵ = $ϵ0*10^E$ = $(@sprintf(\"%.5g\",ϵ))"
end

# ╔═╡ 1b0d0069-d7f1-45f9-8690-d68655340f49
begin
scatter(n0:n0+N,a.(n0:n0+N),markersize=0.7,legend=false)
hline!([L],width=1)
hline!([L-ϵ,L+ϵ])
end

# ╔═╡ Cell order:
# ╟─ea8cc0e4-9920-11eb-360c-a1ca0e5fea2a
# ╟─483ce98d-3409-4fed-a360-89291d643ae1
# ╟─12f9cc2e-e335-4242-89ac-e3819c944266
# ╟─a1b36581-d9fc-4b41-8c41-1a6d64d95d4a
# ╠═742547cf-890a-4b6f-8d24-92772b5b3f14
# ╟─3deee887-5ec0-4b01-84bb-b1a4edd0cffb
# ╟─52c81bda-00c9-4c68-ab02-8a78dd8e79a0
# ╟─35838740-5278-4e35-bc1f-58e783789101
# ╟─13fce500-346f-4446-a6b5-9c5270c25d6f
# ╟─1b0d0069-d7f1-45f9-8690-d68655340f49


"""
`ScaledKernel(k::Kernel,σ²::Real)`
Return a kernel premultiplied by the variance `σ²` : `σ² k(x,x')`
"""
struct ScaledKernel{Tk<:Kernel, Tσ²<:Real} <: Kernel
    kernel::Tk
    σ²::Vector{Tσ²}
end

function ScaledKernel(kernel::Tk, σ²::Tσ²=1.0) where {Tk<:Kernel,Tσ²<:Real}
    @check_args(ScaledKernel, σ², σ² > zero(Tσ²), "σ² > 0")
    return ScaledKernel{Tk, Tσ²}(kernel, [σ²])
end

kappa(k::ScaledKernel, x) = first(k.σ²) * kappa(k.kernel, x)

metric(k::ScaledKernel) = metric(k.kernel)

Base.:*(w::Real, k::Kernel) = ScaledKernel(k, w)

Base.show(io::IO, κ::ScaledKernel) = printshifted(io, κ, 0)

function printshifted(io::IO, κ::ScaledKernel, shift::Int)
    printshifted(io, κ.kernel, shift)
    print(io,"\n" * ("\t"^(shift+1)) * "- σ² = $(first(κ.σ²))")
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.10.3
#   kernelspec:
#     display_name: Julia 1.6.3
#     language: julia
#     name: julia-1.6
# ---

# %%
using AbstractAlgebra
AbstractAlgebra.PrettyPrinting.set_html_as_latex(true)

using AbstractAlgebra
AbstractAlgebra.PrettyPrinting.set_html_as_latex(true)

eqstr(x) = "\\ \\text{is " * string(x) * "}"
eqstr(x::SetElem) = "=" * sprint(x) do io, x; show(io, "text/latex", x) end

function dispeq(lhs, rhs)
    latexstr = "\$\\displaystyle $lhs " * eqstr(rhs) * "\$"
    display("text/html", latexstr)
end

function euclidean(f, g, stop = 1)
    while true
        r = mod(f, g)
        iszero(r) && return g
        degree(r) ≤ stop && return r
        f, g = g, r
    end
end

function rootdeg1(f)
    C = coefficients(f) |> collect
    -first(C)//last(C)
end

function commonroot(f, g)
    m = euclidean(f, g, 1)
    rootdeg1(m)
end

function revpoly(z, x, f)
    m = degree(f)
    C = coefficients(f)
    sum(a*z^(i-1)*x^(m-i+1) for (i, a) in enumerate(C))
end

function val(h::FracElem, θ, L)
    n = L(numerator(h)(θ))
    d = L(denominator(h)(θ))
    n//d
end

function testval(h::FracElem, θ, β, L)
    n = L(numerator(h)(θ))
    d = L(denominator(h)(θ))
    n == β*d
end

function calcallresults(z, x, f, g)
    f_plus = (-1)^degree(f)*f(z - x)
    F_plus = resultant(f_plus, g)
    β_plus = commonroot(g, f_plus)

    f_mult = (-1)^degree(f)*revpoly(z, x, f)
    F_mult = resultant(f_mult, g)
    β_mult = commonroot(g, f_mult)

    K = FractionField(base_ring(z))
    B1, α = K["α"]
    L1 = ResidueField(B1, numerator(f(z))(α))
    B2, β = L1["β"]
    L = ResidueField(B2, numerator(g(z))(β))

    val_F_plus = val(F_plus, α + β, L)
    test_β_plus = testval(β_plus, α + β, β, L)
    val_F_mult = val(F_mult, α * β, L)
    test_β_mult = testval(β_mult, α * β, β, L)

    S_plus = subresultant((-1)^degree(f)*f(z-x), g, 1)
    rootS_plus = rootdeg1(S_plus)

    F_plus, β_plus, F_mult, β_mult, val_F_plus, test_β_plus, val_F_mult, test_β_mult, S_plus, rootS_plus
end

function dispallresults(z, x, f, g)
    F_plus, β_plus, F_mult, β_mult, val_F_plus, test_β_plus, val_F_mult, test_β_mult, S_plus, rootS_plus =
        @time calcallresults(z, x, f, g)
    flush(stdout)
    dispeq("F_\\alpha(x)", f)
    dispeq("F_\\beta(x)", g)
    dispeq("R_{\\alpha + \\beta}(z)", F_plus)
    dispeq("R_{\\alpha + \\beta}(\\alpha + \\beta)", val_F_plus)
    dispeq("\\beta_\\mathrm{plus}(z)", β_plus)
    dispeq("\\beta_\\mathrm{plus}(\\alpha + \\beta) = \\beta", test_β_plus)
    dispeq("R_{\\alpha\\beta}(z)", F_mult)
    dispeq("R_{\\alpha\\beta}(\\alpha\\beta)", val_F_mult)
    dispeq("\\beta_\\mathrm{mult}(z)", β_mult)
    dispeq("\\beta_\\mathrm{mult}(\\alpha\\beta) = \\beta", test_β_mult)
    dispeq("\\text{1-subresultant}", S_plus)
    dispeq("\\text{root of 1-subresultant}", rootS_plus)
end

safecoeff(f, k) = 0 ≤ k ≤ degree(f) ? coeff(f, k) : zero(base_ring(f))

function subresultant_matrix(p::PolyElem{T}, q::PolyElem{T}, k) where T <: RingElement
    check_parent(p, q)
    R = parent(p)
    m = degree(p)
    n = degree(q)
    if length(p) == 0 || length(q) == k || m + n < 2k
       return zero_matrix(R, 0, 0)
    end
    M = zero_matrix(R, m + n - 2k, m + n - 2k)
    x = gen(R)
    for i in 1:n-k
        for j in 1:m+n-2k-1
           M[i, j] = safecoeff(p, m + (i-1) - (j-1))
        end
        M[i, end] = x^(n-k-i)*p
    end
    for i in 1:m-k
        for j in 1:m+n-2k-1
           M[n-k+i, j] = safecoeff(q, n + (i-1) - (j-1))
        end
        M[n-k+i, end] = x^(m-k-i)*q
    end
    return M
end

subresultant(p, q, k) = det(subresultant_matrix(p, q, k))

# %%
R, (a, b, c, p, q, r, s) = ZZ["a", "b", "c", "p", "q", "r", "s"]
K = FractionField(R)
Rz, z = R["z"]
Kz = FractionField(Rz)
Rx, x = Kz["x"]

# %%
f = x^3 + a*x^2 + b*x + c
g = x^4 + p*x^3 + q*x^2 + r*x + s
subresultant_matrix(f, g, 0)

# %%
S0 = sylvester_matrix(f, g)

# %%
subresultant(f, g, 0) == resultant(f, g) == det(S0)

# %% tags=[]
subresultant_matrix(f, g, 1)

# %%
S1 = Rx[
    1 a b c 0
    0 1 a b c*x
    0 0 1 a b*x+c
    1 p q r s*x
    0 1 p q r*x+s
]

# %%
subresultant(f, g, 1)

# %% tags=[]
subresultant(f, g, 1) == det(S1)

# %%
S2 = Rx[
    1 a b*x^2+c*x
    0 1 a*x^2+b*x+c
    1 p q*x^2+r*x+s
]

# %%
subresultant(f, g, 2)

# %%
subresultant(f, g, 2) == det(S2)

# %%
f = x^3 - a
g = x^4 - p
dispallresults(z, x, f, g)

# %%
subresultant_matrix(-f(z-x), g, 1)

# %%
s1 = Rx[
    1 -3z 3z^2 -z^3+a 0
    0 1   -3z  3z^2   (-z^3+a)*x
    0 0   1    -3z    3z^2*x-z^3+a
    1 0   0    0      -p*x
    0 1   0    0      -p
]

# %%
S_plus = subresultant(-f(z-x), g, 1)

# %%
subresultant(-f(z-x), g, 1) == det(s1)

# %%
rootS_plus = rootdeg1(S_plus)
dispeq("\\text{root of 1-subresultant}", rootS_plus)

# %% tags=[]
f = x^3 - a
g = x^3 + p*x + q
dispallresults(z, x, f, g)

# %% tags=[]
f = x^3 - a
g = x^3 + p*x + q
dispallresults(z, x, g, f)

# %%
h1 = g(z-x) + f

# %%
(3z)^2*f - 3z*x*h1

# %%
h2 = (3z)^2*f - (3z*x + 3z^2 + p)*h1

# %%
subresultant(f, g(z-x), 1)

# %%
f = x^2 - a
g = x^2 - p
dispallresults(z, x, f, g)

# %%
f = x^2 - a
g = x^3 - p
dispallresults(z, x, f, g)

# %%
f = x^2 - a
g = x^4 - p
dispallresults(z, x, f, g)

# %%
f = x^2 - a
g = x^5 - p
dispallresults(z, x, f, g)

# %%
f = x^3 - a
g = x^3 - p
dispallresults(z, x, f, g)

# %%
f = x^3 - a
g = x^4 - p
dispallresults(z, x, f, g)

# %%
f = x^3 - a
g = x^5 - p
dispallresults(z, x, f, g)

# %%
f = x^3 + a*x + b
g = x^3 + p*x + q
dispallresults(z, x, f, g)

# %%
f = x^2 - a
g = x^3 + p*x^2 + q*x + r
dispallresults(z, x, f, g)

# %%
f = x^3 + a*x + b
g = x^4 + p*x^2 + q*x + r
dispallresults(z, x, f, g)

# %%


# test_bb_operations()
A = BoundingBox(10,10,30,20)
B = BoundingBox(5,20,20,10)
C = A+B
D = BoundingBox(5,10,35,20)
@test C == D

# @test_throws InexactError BoundingBox(0,0,-1,-1)

# test_find_mesh_bb()
  dst = [4.0 2.0;
          8.0 2.0;
          6.0 10.0]
  bounds = find_mesh_bb(dst)
  @test bounds == BoundingBox(4,2,4,8)

  dst = [0.0 0.0;
          0.0 0.0;
          0.0 0.0]
  bounds = find_mesh_bb(dst)
  @test bounds == BoundingBox(0,0,0,0)

# test_tform_bb()
tform = [1 0 0;
        0 1 0;
        0 0 1];
bb = BoundingBox(100,100,100,100)
tbb = tform_bb(bb, tform)
@test bb == tbb

tform = [1 0 0;
        0 1 0;
        100 100 1];
bb = BoundingBox(100,100,100,100)
tbb = tform_bb(bb, tform)
@test BoundingBox(200,200,100,100) == tbb

sz = (100, 100)
tform = [cos(pi/6) -sin(pi/6) 0;
        sin(pi/6) cos(pi/6) 0;
        0 0 1];
bb = BoundingBox(0, -50, 136.60254037844388, 136.60254037844388)
tbb = tform_bb(sz_to_bb(sz), tform)
@test_approx_eq bb_to_pts(bb) bb_to_pts(tbb)

# test_snap_bb()
bb = snap_bb(sz_to_bb((100.5, 200.75)))
tbb = sz_to_bb((101, 201))
@test bb_to_pts(bb) == bb_to_pts(tbb)

bb = snap_bb(sz_to_bb((100.9, 200.1)))
tbb = sz_to_bb((101, 201))
@test bb_to_pts(bb) == bb_to_pts(tbb)

bb = snap_bb(BoundingBox(100.9, 200.1, 100.5, 100.2))
tbb = sz_to_bb((101, 201))
tbb = BoundingBox(100, 200, 102, 101)
@test bb_to_pts(bb) == bb_to_pts(tbb)

a = BoundingBox(0,0,10,10)
b = BoundingBox(0,0,5,5)
@test a+b == BoundingBox(0,0,10,10)
@test a-b == BoundingBox(0,0,5,5)

a = BoundingBox(2,2,10,10)
b = BoundingBox(0,0,5,5)
@test a+b == BoundingBox(0,0,12,12)
@test a-b == BoundingBox(2,2,3,3)

a = BoundingBox(20,20,10,10)
b = BoundingBox(0,0,5,5)
@test a+b == BoundingBox(0,0,30,30)
@test isnan((a-b).i)


"""
    solve

Solve equation.
"""
function solve end

include("solve_advection.jl")
include("solve_diffusion.jl")
include("solve_burgers.jl")