include("utils.jl")

@testset "Grassmann" begin
    @testset "Real" begin
        M = Grassmann(3,2)
        @testset "Basics" begin
            @test repr(M) == "Grassmann(3, 2, ℝ)"
            @test representation_size(M) == (3,2)
            @test manifold_dimension(M) == 2
            @test !is_manifold_point(M,[1., 0., 0., 0.])
            @test !is_tangent_vector(M, [1.0 0.0; 0.0 1.0; 0.0 0.0], [0., 0., 1., 0.])
            @test_throws DomainError is_manifold_point(M, [2. 0.; 0. 1.; 0. 0.],true)
            @test_throws DomainError is_tangent_vector(M, [2. 0.; 0. 1.; 0. 0.],zeros(3,2),true)
            @test_throws DomainError is_tangent_vector(M, [1. 0.; 0. 1.; 0. 0.],ones(3,2),true)
            @test is_manifold_point(M, [1.0 0.0; 0.0 1.0; 0.0 0.0], true)
            @test_throws DomainError is_manifold_point(M, 1im*[1.0 0.0; 0.0 1.0; 0.0 0.0], true)
            @test is_tangent_vector(M, [1.0 0.0; 0.0 1.0; 0.0 0.0], zero_tangent_vector(M,[1.0 0.0; 0.0 1.0; 0.0 0.0]), true)
            @test_throws DomainError is_tangent_vector(M, [1.0 0.0; 0.0 1.0; 0.0 0.0], 1im*zero_tangent_vector(M,[1.0 0.0; 0.0 1.0; 0.0 0.0]), true)
            @test injectivity_radius(M) == π/2
            @test injectivity_radius(M, ExponentialRetraction()) == π/2
            @test injectivity_radius(M, [1.0 0.0; 0.0 1.0; 0.0 0.0]) == π/2
            @test injectivity_radius(M, [1.0 0.0; 0.0 1.0; 0.0 0.0], ExponentialRetraction()) == π/2
        end
        types = [Matrix{Float64}, ]
        TEST_STATIC_SIZED && push!(types, MMatrix{3, 2, Float64})

        TEST_FLOAT32 && push!(types, Matrix{Float32})
        basis_types = (ProjectedOrthonormalBasis(:gram_schmidt),)
        @testset "Type $T" for T in types
            x = [1.0 0.0; 0.0 1.0; 0.0 0.0]
            v = [0.0 0.0; 0.0 0.0; 0.0 1.0]
            y = exp(M,x,v)
            w = [0.0 1.0; -1.0 0.0; 1.0 0.0]
            z = exp(M,x,w)
            pts = convert.(T, [x,y,z])
            test_manifold(
                M,
                pts,
                test_exp_log = true,
                test_injectivity_radius = false,
                test_project_tangent = true,
                test_vector_transport = false,
                test_forward_diff = false,
                test_reverse_diff = false,
                test_vee_hat = false,
                retraction_methods = [PolarRetraction(), QRRetraction()],
                inverse_retraction_methods = [PolarInverseRetraction(), QRInverseRetraction()],
                #basis_types_vecs = basis_types,
                exp_log_atol_multiplier = 10.0,
                is_tangent_atol_multiplier = 10.0,
            )

            @testset "inner/norm" begin
                v1 = inverse_retract(M, pts[1], pts[2], PolarInverseRetraction())
                v2 = inverse_retract(M, pts[1], pts[3], PolarInverseRetraction())

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

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

    @testset "Complex" begin
        M = Grassmann(3,2,ℂ)
        @testset "Basics" begin
            @test repr(M) == "Grassmann(3, 2, ℂ)"
            @test representation_size(M) == (3,2)
            @test manifold_dimension(M) == 4
            @test !is_manifold_point(M,[1., 0., 0., 0.])
            @test !is_tangent_vector(M, [1.0 0.0; 0.0 1.0; 0.0 0.0], [0., 0., 1., 0.])
            @test Manifolds.allocation_promotion_function(M,exp!,(1,)) == complex
            @test_throws DomainError is_manifold_point(M, [2. 0.; 0. 1.; 0. 0.],true)
            @test_throws DomainError is_tangent_vector(M, [2. 0.; 0. 1.; 0. 0.],zeros(3,2),true)
            @test_throws DomainError is_tangent_vector(M, [1. 0.; 0. 1.; 0. 0.],ones(3,2),true)
            @test is_tangent_vector(M, [1.0 0.0; 0.0 1.0; 0.0 0.0], 1im*zero_tangent_vector(M,[1.0 0.0; 0.0 1.0; 0.0 0.0]))
            @test is_manifold_point(M, [1.0 0.0; 0.0 1.0; 0.0 0.0])
            @test injectivity_radius(M) == π/2
        end
        types = [
            Matrix{ComplexF64},
        ]
        @testset "Type $T" for T in types
            x = [0.5+0.5im 0.5+0.5im; 0.5+0.5im -0.5-0.5im; 0.0 0.0]
            v = [0.0 0.0; 0.0 0.0; 0.0 1.0]
            y = exp(M,x,v)
            w = [0.0 1.0; -1.0 0.0; 1.0 0.0]
            z = exp(M,x,w)
            pts = convert.(T, [x,y,z])
            test_manifold(
                M,
                pts,
                test_exp_log = true,
                test_injectivity_radius = false,
                test_project_tangent = true,
                test_vector_transport = false,
                test_forward_diff = false,
                test_reverse_diff = false,
                test_vee_hat = false,
                retraction_methods = [PolarRetraction(), QRRetraction()],
                inverse_retraction_methods = [PolarInverseRetraction(), QRInverseRetraction()],
                exp_log_atol_multiplier = 10.0,
                is_tangent_atol_multiplier = 10.0,
            )

            @testset "inner/norm" begin
                v1 = inverse_retract(M, pts[1], pts[2], PolarInverseRetraction())
                v2 = inverse_retract(M, pts[1], pts[3], PolarInverseRetraction())

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

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


using FractionalCalculus
using Plots

x=collect(0.1:0.1:5)
y=collect(0.1:0.1:5)
z=zeros(50, 50)

for j in range(1, 50, step=1)
    z[j, :] = fracdiff(sin, 0.02*j, y, 0.01, RLDiff_Approx())
end

plot(x, y, z, st=:surface)

savefig("./3dexample.png")


using Test
using LinearAlgebra
using QuantumPropagators

@testset "random Hermitian state" begin

    # Input
    N = 1000

    X = rand(ComplexF64, (N, N))
    H = Hermitian(X)

    dt = 0.5

    Ψ₀ = rand(ComplexF64, N)
    Ψ₀ ./= norm(Ψ₀)

    @test norm(Ψ₀) ≈ 1

    # Expected
    U = exp(-im * H * dt)
    Ψ_out_expected = U * Ψ₀
    @test norm(Ψ_out_expected) ≈ 1

    Ψ = copy(Ψ₀)
    wrk = ExpPropWrk(Ψ₀)
    expprop!(Ψ, H, dt, wrk)
    Ψ_out = copy(Ψ)
    @test norm(Ψ_out) ≈ 1

    # Comparison (should be exact to machine precision)
    @test norm(Ψ_out - Ψ_out_expected) ≈ 0
end


# Test chordal extension
workspace()
include("qdldl.jl")
include("/Users/Micha/Dropbox/Research/OSSDP/Code/src/Graph.jl")

using QDLDL,GraphModule, Base.Test

rng = Random.MersenneTwister(131123)

numProblems = 10

# g = Graph(A,false)
# println("Original Graph is chordal?: $(GraphModule.isPerfectOrdering(g))")
# # A = A  +(2*5+1)*speye(5);
# F = QDLDL.qdldl(A,logical=true,perm=nothing)
# L = F.L

# C = full(L+L') + eye(5)

# g_c = Graph(sparse(C))
# println("Extended Graph is chordal?: $(GraphModule.isPerfectOrdering(g_c))")

@testset "QDLDL routine for chordal embedding" begin

  for iii=1:numProblems

    dim = rand(rng,50:1:200);
    density=rand(rng,0.1:0.1:0.4)
    A = sprand(rng,dim,dim,density)

    A = 0.5*(A + A')
    A = A + (2*dim+1)*speye(dim);

    g = Graph(A,false)

    if (GraphModule.isPerfectOrdering(g) == true)
      continue
    else
      F = QDLDL.qdldl(A,logical=true,perm=nothing)
      L = F.L
      C = L+L' + speye(dim)
      g_c = Graph(C)
    end

    # check results
    @test GraphModule.isPerfectOrdering(g_c)

  end
end


# calculate the stoichiometry coefficients
function stoich_coeff(species::Vector{T}, reaction::DiffEqBiological.ReactionStruct, e_order::Vector{Int}) where T <: AlgebraSet
    coeff = 1
    for ind in eachindex(e_order)
        coeff *= DiffEqBiological.get_stoch_diff(reaction, Symbol(species[ind]))^e_order[ind]
    end
    coeff
end

# calculate the sign coefficients
function sign_coeff(a_ords::Vector{Int}, b_ords::Vector{Int})
    *((-1).^(a_ords .- b_ords)...)
end

# calculate the exponentiated mean terms
function α_coeff(raw_means::Vector{T}, n_orders::Vector{Int}, k_orders::Vector{Int}) where T <: AlgebraSet
    *(raw_means.^(n_orders .- k_orders)...)
end

# calculate the bionomial coefficients
function binomial_coeff(a_ords::Vector{Int}, b_ords::Vector{Int})
    *(binomial.(a_ords, b_ords)...)
end

# calculate the 1 / n! coefficients
function inv_factorial_calc(orders::Vector{Int})
    1 / *(factorial.(orders)...)
end

# to rewrite the mass rate calculation of DiffEqBiological, so that it can be differentiated later on.
function mass_rate_MA(reaction::DiffEqBiological.ReactionStruct; algebra::Algebra=sympy())
    #println(SymPy.Sym(string(reaction.rate_org)))
    rate = algebra.Sym(string(reaction.rate_org))
    if reaction.is_pure_mass_action
        for sub in reaction.substrates
             for i in 1:sub.stoichiometry
                rate *= (algebra.Sym(sub.reactant) + 1 - i) / i
            end
        end
    end
    rate
end

# convert raw moments to central moments
function raw_to_cen(momentstruct::MomentStruct, n_order::Vector{Int})
    eq = zero(eltype(momentstruct.species))
    for ord in Iterators.product([0:n_max for n_max in n_order]...)
        __n =collect(ord)
        eq += binomial_coeff(n_order, __n) * α_coeff(momentstruct.species, n_order, __n) * momentstruct.cen_moments[__n]
    end
    eq
end

# convert central moments to raw moments (currently unused)
function cen_to_raw(momentstruct::MomentStruct, n_order::Vector{Int})
	eq = zero(eltype(momentstruct.species))
	for ord in Iterators.product([0:n_max for n_max in n_order]...)
		__n =collect(ord)
		if sum(__n) == 1
			eq += binomial_coeff(n_order, __n) * sign_coeff(n_order, __n) * α_coeff(momentstruct.species, n_order, __n) * momentstruct.species[findall(isone,__n)[1]]
		else
			eq += binomial_coeff(n_order, __n) * sign_coeff(n_order, __n) * α_coeff(momentstruct.species, n_order, __n) * momentstruct.raw_moments[__n]
		end
	end
	eq
end

# compute taylor expansion
function taylor_expand(expr::AlgebraSet, momentstruct::MomentStruct; algebra::Algebra=sympy())
    eq = zero(algebra)
    for orders in keys(momentstruct.cen_moments)
        diff_args = reduce(append!, repeat([momentstruct.species[i]], orders[i]) for i in eachindex(orders))
        factorial_coeff = inv_factorial_calc(orders)
        eq += foldl(algebra.diff, (expr, diff_args...))*momentstruct.cen_moments[orders]*factorial_coeff
    end
    eq
end

# generate ode expressions and function from symbolic arrays
function get_ode(network::DiffEqBase.AbstractReactionNetwork, moments::MomentStruct, dμ::Vector{T}, dM::Vector{T}) where T <: AlgebraSet
    filtered_cen_moments = filter(p->sum(p.first)>1, moments.cen_moments)
    central_dict = OrderedDict{Symbol, Int}(Symbol(pair[2])=>i+moments.num_species for (i,pair) in enumerate(filtered_cen_moments))
    ode_expr = Vector{Expr}(undef, moments.num_species + length(central_dict))
    for i = 1:moments.num_species
        ex = dμ[i]
        ode_expr[i] = :(internal_var___du[$i] = $(Meta.parse(string(ex))))
    end
    for i = (moments.num_species+1):(moments.num_species+length(central_dict))
        ex = dM[i-moments.num_species]
        ode_expr[i] = :(internal_var___du[$i] = $(Meta.parse(string(ex))))
    end
    combined_dict = merge(speciesmap(network), central_dict)
    ode_func = DiffEqBiological.make_func(ode_expr, combined_dict, paramsmap(network))
    ode_expr, ODEFunction(eval(ode_func))
end


"""
    Equilibrium!(feq, height,velocityx, velocityy, vsquare, gravity)

Calculation of the equilibrium distribution `feq` for the shallow water lattice Boltzmann method.

# Arguments

- `feq :: Array{<:Number,3}`: Equilibrium distribution function, to be calculated
- `height :: Array{<:Number,2}`: The height field `` h(\\mathbf{x},t)``
- `velocityx :: Array{<:Number,2}`: x-component of the macroscopic velocity
- `velocityy :: Array{<:Number,2}`: y-component of the macroscopic velocity
- `vsquare :: Array{<:Number,2}`: Dummy array that is preallocated to be filled with the square of the velocity vector
- `gravity <: Number`: Strength of the gravitational acceleration in lattice units

# Mathematics

The detailed motivation and derivation can be found in the article of Salmon.
Similar to the standard Navier-Stokes approximating lattice Boltzmann methods the equilibrium distribution `feq`
is an expansion to second order in velocity `u`.
If you want it is an slow mode part of the shallow water theory and thus the equilibrium is given as

`` f_i^{\\text{eq}} = h \\bigg(1 - \\frac{5}{6}g h - \\frac{2}{3}u^2\\bigg),\\quad i=0 \\newline
   f_i^{\\text{eq}} = w_i h \\bigg(g h + 3 \\mathbf{c}_i\\cdot\\mathbf{u} + \\frac{9}{2}(\\mathbf{c}_i\\cdot\\mathbf{u})^2) + \\frac{3}{2}u^2\\bigg),\\quad else ``

where ``g`` is the gravitational acceleration (in lattice units) and ``w_i, \\mathbf{c}_i`` are the weights and lattice velocities.

It has been shown that it is possible to get rid of the gravity driven term in the equilibrium distribution, thus

`` f_i^{\\text{eq}} = h \\bigg(1 - \\frac{2}{3}u^2\\bigg),\\quad i=0 \\newline
f_i^{\\text{eq}} = w_i h \\bigg(3 \\mathbf{c}_i\\cdot\\mathbf{u} + \\frac{9}{2}(\\mathbf{c}_i\\cdot\\mathbf{u})^2) + \\frac{3}{2}u^2\\bigg),\\quad else ``

then of course the topography gradient has to be included as a force term.

# Examples
```jldoctest
julia> using Swalbe, Test

julia> feq = zeros(5,5,9); ρ = ones(5,5); ux = fill(0.1,5,5); uy = zeros(5,5);

julia> Swalbe.equilibrium!(feq, ρ, ux, uy, zeros(5,5), 0.1) # Supply dummy u^2 as well.

julia> feq[:,:,1]
5×5 Matrix{Float64}:
 0.91  0.91  0.91  0.91  0.91
 0.91  0.91  0.91  0.91  0.91
 0.91  0.91  0.91  0.91  0.91
 0.91  0.91  0.91  0.91  0.91
 0.91  0.91  0.91  0.91  0.91

julia> Swalbe.equilibrium!(feq, ρ, ux, uy, zeros(5,5)) # Use the dispatch without g.

julia> @test all(feq[:,:,1] .≈ 1 - 2/3 * 0.01)
Test Passed
```

# References

- [Salmon](https://www.ingentaconnect.com/contentone/jmr/jmr/1999/00000057/00000003/art00005#)
- [Dellar](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.65.036309)
- [Peng et al.](https://onlinelibrary.wiley.com/doi/full/10.1002/fld.4726)

"""
function equilibrium!(feq, height, velocityx, velocityy, vsquare, gravity)
    # Views help to circumvent having a loop, which sucks on the GPU
    f0, f1, f2, f3, f4, f5, f6, f7, f8 = viewdists(feq)
    # Some constants, gravity and weights
    g0 = 1.5 * gravity
    w1 = 1/9
    w5 = 1/36

    vsquare .= velocityx .* velocityx .+ velocityy .* velocityy

    # Zeroth dist
    f0 .= height .* (1 .- 5/6 .* gravity .* height .- 2/3 .* vsquare)
    # First
    f1 .= w1 .* height .* (g0 .* height .+ 3 .* velocityx .+ 4.5 .* velocityx.^2 .- 1.5 .* vsquare)
    # Second
    f2 .= w1 .* height .* (g0 .* height .+ 3 .* velocityy .+ 4.5 .* velocityy.^2 .- 1.5 .* vsquare)
    # Third
    f3 .= w1 .* height .* (g0 .* height .- 3 .* velocityx .+ 4.5 .* velocityx.^2 .- 1.5 .* vsquare)
    # Forth
    f4 .= w1 .* height .* (g0 .* height .- 3 .* velocityy .+ 4.5 .* velocityy.^2 .- 1.5 .* vsquare)
    # Fifth
    f5 .= w5 .* height .* (g0 .* height .+ 3 .* (velocityx .+ velocityy) .+
                         4.5 .* (velocityx .+ velocityy).^2 .- 1.5 .* vsquare)
    # Sixth
    f6 .= w5 .* height .* (g0 .* height .+ 3 .* (velocityy .- velocityx) .+
                         4.5 .* (velocityy .- velocityx).^2 .- 1.5 .* vsquare)
    # Seventh
    f7 .= w5 .* height .* (g0 .* height .- 3 .* (velocityx .+ velocityy) .+
                         4.5 .* (velocityx .+ velocityy).^2 .- 1.5 .* vsquare)
    # Eigth
    f8 .= w5 .* height .* (g0 .* height .+ 3 .* (velocityx .- velocityy) .+
                         4.5 .* (velocityx .- velocityy).^2 .- 1.5 .* vsquare)
    return nothing
end

# Dispatch without gravity input
function equilibrium!(feq, height, velocityx, velocityy, vsquare)
    # Views help to circumvent having a loop, which sucks on the GPU
    f0, f1, f2, f3, f4, f5, f6, f7, f8 = viewdists(feq)
    # Some constants, gravity and weights
    w1 = 1/9
    w5 = 1/36

    vsquare .= velocityx .* velocityx .+ velocityy .* velocityy

    # Zeroth dist
    f0 .= height .* (1 .- 2/3 .* vsquare)
    # First
    f1 .= w1 .* height .* (3 .* velocityx .+ 4.5 .* velocityx.^2 .- 1.5 .* vsquare)
    # Second
    f2 .= w1 .* height .* (3 .* velocityy .+ 4.5 .* velocityy.^2 .- 1.5 .* vsquare)
    # Third
    f3 .= w1 .* height .* (-3 .* velocityx .+ 4.5 .* velocityx.^2 .- 1.5 .* vsquare)
    # Forth
    f4 .= w1 .* height .* (-3 .* velocityy .+ 4.5 .* velocityy.^2 .- 1.5 .* vsquare)
    # Fifth
    f5 .= w5 .* height .* (3 .* (velocityx .+ velocityy) .+
                         4.5 .* (velocityx .+ velocityy).^2 .- 1.5 .* vsquare)
    # Sixth
    f6 .= w5 .* height .* (3 .* (velocityy .- velocityx) .+
                         4.5 .* (velocityy .- velocityx).^2 .- 1.5 .* vsquare)
    # Seventh
    f7 .= w5 .* height .* (-3 .* (velocityx .+ velocityy) .+
                         4.5 .* (velocityx .+ velocityy).^2 .- 1.5 .* vsquare)
    # Eigth
    f8 .= w5 .* height .* (3 .* (velocityx .- velocityy) .+
                         4.5 .* (velocityx .- velocityy).^2 .- 1.5 .* vsquare)
    return nothing
end

function equilibrium!(feq, height, velocity, gravity)
    # Views help to circumvent having a loop, which sucks on the GPU
    f0, f1, f2 = viewdists_1D(feq)

    # Zeroth dist
    f0 .= height .* (1 .- 0.5 .* gravity .* height .- velocity.^2)
    # First
    f1 .= height .* (0.25 .* gravity .* height .+ 0.5 .* velocity .+ 0.5 .* velocity.^2)
    # Second
    f2 .= height .* (0.25 .* gravity .* height .- 0.5 .* velocity .+ 0.5 .* velocity.^2)

    return nothing
end

function equilibrium!(feq, height, velocity)
    # Views help to circumvent having a loop, which sucks on the GPU
    f0, f1, f2 = viewdists_1D(feq)

    # Zeroth dist
    f0 .= height .* (1 .- velocity.^2)
    # First
    f1 .= height .* (0.5 .* velocity .+ 0.5 .* velocity.^2)
    # Second
    f2 .= height .* (-0.5 .* velocity .+ 0.5 .* velocity.^2)
    return nothing
end


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:2] <= 1)

                     @variable(m, -3.611664788975792 <= q <= 3.6444325166036777)

                     add_NL_constraint(m, :(1/(1 + exp(-(-0.4383440965768486 + -0.2075365332338559*$(x[1]) + 0.46530061880514095*$(x[2])))) + 1/(1 + exp(-(0.22314644246023452 + -0.6314968564820513*$(x[1]) + 0.5512122876813113*$(x[2])))) + 1/(1 + exp(-(0.30053153463966176 + -0.15357242883325783*$(x[1]) + 0.254162988138138*$(x[2])))) + 1/(1 + exp(-(-0.06895001670910483 + 0.7534492827978192*$(x[1]) + -0.6113176568181604*$(x[2])))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


using UncertaintyQuantification

l = Parameter(1.8, :l) # length
b = Parameter(0.12, :b) # width

h = RandomVariable(Normal(0.24, 0.01), :h) # height

μ = log(10e9^2 / sqrt(1.6e9^2 + 10e9^2))
σ = sqrt(log(1.6e9^2 / 10e9^2 + 1))
E = RandomVariable(LogNormal(μ, σ), :E) # young's modulus

μ = log(5000^2 / sqrt(400^2 + 5000^2))
σ = sqrt(log(400^2 / 5000^2 + 1))
P = RandomVariable(LogNormal(μ, σ), :P) # tip load

μ = log(600^2 / sqrt(140^2 + 600^2))
σ = sqrt(log(140^2 / 600^2 + 1))
ρ = RandomVariable(LogNormal(μ, σ), :ρ) # density

c = GaussianCopula([1 0.8; 0.8 1])
jd = JointDistribution([E ρ], c)

inputs = [l, b, h, P, jd]

inertia = Model(df -> df.b .* df.h .^ 3 / 12, :I)

displacement = Model(
    df ->
        (df.ρ .* 9.81 .* df.b .* df.h .* df.l .^ 4) ./ (8 .* df.E .* df.I) .+
        (df.P .* df.l .^ 3) ./ (3 .* df.E .* df.I),
    :w,
)

max_displacement = 0.01

# Compute probability of failure using standard Monte Carlo
mc = MonteCarlo(10^6)

mc_pf, mc_samples = probability_of_failure(
    [inertia, displacement], df -> max_displacement .- df.w, inputs, mc
)

println(
    "Monte Carlo probability of failure $mc_pf ($(size(mc_samples, 1)) model evaluations)"
)

# Compute probability of failure using Line Sampling
ls = LineSampling(50)

ls_pf, ls_samples = probability_of_failure(
    [inertia, displacement], df -> max_displacement .- df.w, inputs, ls
)

println(
    "Line Sampling probability of failure: $ls_pf ($(size(ls_samples, 1)) model evaluations)",
)

subset = UncertaintyQuantification.SubSetSimulation(2000, 0.1, 10, Uniform(-0.5, 0.5))

subset_pf, subset_samples = probability_of_failure(
    [inertia, displacement], df -> max_displacement .- df.w, inputs, subset
)

println(
    "Subset Simulation probability of failure: $subset_pf ($(size(subset_samples, 1)) model evaluations)",
)


function main()
	N = 600_851_475_143
	for i = isqrt(N):-1:1
		if 0 == N % i && isprime(i)
			println(i)
			break
		end
	end
end

main()


#=

Example demonstrating sparsity control by means of the "K"-factor passed to the lazy AFW variant

A larger K >= 1 favors sparsity by favoring optimization over the current active set rather than
adding a new FW vertex.

The default for AFW is K = 2.0

=#

using FrankWolfe
using ProgressMeter
using Arpack
using Plots
using DoubleFloats
using ReverseDiff

using LinearAlgebra
using Random

n = Int(1e4)
k = 1000

s = rand(1:100)
@info "Seed $s"
Random.seed!(s)

xpi = rand(n);
total = sum(xpi);

# here the optimal solution lies in the interior if you want an optimal solution on a face and not the interior use:
# const xp = xpi;

const xp = xpi ./ total;

f(x) = norm(x - xp)^2
function grad!(storage, x)
    @. storage = 2 * (x - xp)
end

const lmo = FrankWolfe.KSparseLMO(100, 1.0)

## other LMOs to try
# lmo_big = FrankWolfe.KSparseLMO(100, big"1.0")
# lmo = FrankWolfe.LpNormLMO{Float64,5}(1.0)
# lmo = FrankWolfe.ProbabilitySimplexOracle(1.0);
# lmo = FrankWolfe.UnitSimplexOracle(1.0);

const x00 = FrankWolfe.compute_extreme_point(lmo, rand(n))

## example with BirkhoffPolytopeLMO - uses square matrix.
# const lmo = FrankWolfe.BirkhoffPolytopeLMO()
# cost = rand(n, n)
# const x00 = FrankWolfe.compute_extreme_point(lmo, cost)

function build_callback(trajectory_arr)
    return function callback(state)
        return push!(trajectory_arr, (Tuple(state)[1:5]..., length(state.active_set)))
    end
end

FrankWolfe.benchmark_oracles(f, grad!, () -> randn(n), lmo; k=100)

x0 = deepcopy(x00)
@time x, v, primal, dual_gap, trajectory_shortstep = FrankWolfe.frank_wolfe(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.Shortstep(),
    L=2,
    print_iter=k / 10,
    emphasis=FrankWolfe.memory,
    verbose=true,
    trajectory=true,
);

trajectory_adaptive = []
callback = build_callback(trajectory_adaptive)

x0 = deepcopy(x00)
@time x, v, primal, dual_gap, _ = FrankWolfe.away_frank_wolfe(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.Adaptive(),
    print_iter=k / 10,
    emphasis=FrankWolfe.memory,
    verbose=true,
    trajectory=true,
    callback=callback,
);

println("\n==> Lazy AFW.\n")

trajectory_adaptiveLoc15 = []
callback = build_callback(trajectory_adaptiveLoc15)

x0 = deepcopy(x00)
@time x, v, primal, dual_gap, _ = FrankWolfe.away_frank_wolfe(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.Adaptive(),
    print_iter=k / 10,
    emphasis=FrankWolfe.memory,
    verbose=true,
    lazy=true,
    K=1.5,
    trajectory=true,
    callback=callback,
);

trajectory_adaptiveLoc2 = []
callback = build_callback(trajectory_adaptiveLoc2)

x0 = deepcopy(x00)
@time x, v, primal, dual_gap, trajectory_adaptiveLoc2 = FrankWolfe.away_frank_wolfe(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.Adaptive(),
    print_iter=k / 10,
    emphasis=FrankWolfe.memory,
    verbose=true,
    lazy=true,
    K=2.0,
    trajectory=true,
    callback=callback,
);

trajectory_adaptiveLoc4 = []
callback = build_callback(trajectory_adaptiveLoc4)

x0 = deepcopy(x00)
@time x, v, primal, dual_gap, _ = FrankWolfe.away_frank_wolfe(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.Adaptive(),
    print_iter=k / 10,
    emphasis=FrankWolfe.memory,
    verbose=true,
    K=4.0,
    lazy=true,
    trajectory=true,
    callback=callback,
);

trajectory_adaptiveLoc10 = []
callback = build_callback(trajectory_adaptiveLoc10)

x0 = deepcopy(x00)
@time x, v, primal, dual_gap, _ = FrankWolfe.away_frank_wolfe(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.Adaptive(),
    print_iter=k / 10,
    emphasis=FrankWolfe.memory,
    lazy=true,
    K=10.0,
    verbose=true,
    trajectory=true,
    callback=callback,
);

# Reduction primal/dual error vs. sparsity of solution

dataSparsity =
    [trajectory_adaptive, trajectory_adaptiveLoc15, trajectory_adaptiveLoc2, trajectory_adaptiveLoc4, trajectory_adaptiveLoc10]
labelSparsity = ["AFW", "LAFW-K-1.5", "LAFW-K-2.0", "LAFW-K-4.0", "LAFW-K-10.0"]

FrankWolfe.plot_sparsity(dataSparsity, labelSparsity, legend_position=:topright)


BUCH_DISCARDED = 0
BUCH_NF        = 0
BUCH_SIZE      = 0
BUCH_REDUCED   = 0

function zerobuch()
    global BUCH_DISCARDED
    global BUCH_NF
    global BUCH_SIZE
    global BUCH_REDUCED
    BUCH_DISCARDED = 0
    BUCH_NF    = 0
    BUCH_SIZE      = 0
    BUCH_REDUCED   = 0
end

#-----------------------------------------------------------------------------
# Regular reductions

# try to reduce f with g in a regular way and return result
function reduction_step(f, g)
    @info "reducing $(f) with $(g)" f g

    for t in terms(f)
        # if divides
        success, u = divides(t, leading_monomial(g))
        if success
            u = divexact(t, leading_term(g))
            @info "$(leading_term(g)) | $t"

            newf = f - u*g
            @info "reduction! $(f) - $(u)*$(g) --> $newf"
            return true, newf
        end
    end
    return false, f
end

# // same as above, but in terms of set G //
# try to reduce f with G in a regular way and return result
function reduction_step(f, G::Vector{T}) where {T}
    for g in G
        success, newf = reduction_step(f, g)
        if success
            return true, newf
        end
    end
    return false, f
end

# regular normal form of f w.r.t. G
function normal_form(f, G)
    @info "computing normalform of $f w.r.t $G.."

    success = true
    newf = f
    while success
        success, newf = reduction_step(newf, G)
        @info "reduction $success"
    end
    newf
end

#-----------------------------------------------------------------------------
# S-polynomial

# returns u, v,  such that
# u*a - v*b is spoly(a, b)
function mults(a, b)
    u = lcm(leading_monomial(a), leading_monomial(b))
    u = divexact(u, leading_term(a))

    v = lcm(leading_monomial(a), leading_monomial(b))
    v = divexact(v, leading_term(b))

    u, v
end

# S-polynomial of a and b
function spoly(a, b)
    u, v = mults(a, b)
    u*a - v*b
end

function spoly(ab)
    spoly(ab...)
end

#-----------------------------------------------------------------------------

function buchberger_criterion1(f)
    f1, f2 = f
    if isone(gcd(leading_monomial(f1), leading_monomial(f2)))
        @warn "DISCARDED $((f1, f2)) by 1st Criterion"
        global BUCH_DISCARDED
        BUCH_DISCARDED += 1
        return true
    end
    return false
end

#-----------------------------------------------------------------------------
# Groebner basis things

function groebner_basis(F)
    G = deepcopy(F)

    P = []
    for i in 1:length(G)
        for j in i+1:length(G)

            if buchberger_criterion1((G[i], G[j]))
                continue
            end

            push!(P, (G[i], G[j]))
        end
    end

    @info "generated initial G and P:"
    @info "F = $F"
    @info "G = $G"
    @info "P = $P"

    while !isempty(P)
        sort!(P, by=leading_monomial ∘ spoly)
        f = popfirst!(P)
        @info "selected $f"

        global BUCH_NF
        BUCH_NF += 1
        fNF = normal_form(spoly(f), G)

        if iszero(fNF)
            global BUCH_REDUCED
            BUCH_REDUCED += 1
        else
            for fj in G
                if buchberger_criterion1((fNF, fj)) || iszero(spoly(fNF, fj))
                    continue
                end

                push!(P, (fNF, fj))
                @info "$fNF  $fj"
                @warn "SPOLY $(last(P))"
            end

            # update G
            fNFnormalzed = fNF * inv(leading_coefficient(fNF))
            push!(G, fNFnormalzed)
        end

        @info "updated G and P"
        @info "G = $G"
        @info "P = $P"
    end

    global BUCH_SIZE
    BUCH_SIZE = length(G) - length(F)

    global BUCH_REDUCED
    global BUCH_DISCARDED

    G
end

#-----------------------------------------------------------------------------

R, (x,y,z) = PolynomialRing(QQ, ["x","y", "z"], ordering=:degrevlex)

F = [x + y + z, x*y + y*z + x*z, x*y*z - 1]

# evaluation
G = groebner_basis(F)
println("############")
println(G)


using DelaySSAToolkit, Catalyst
using LaTeXStrings
using Plots
theme(:vibrant)

rn = @reaction_network begin
    C, 0 --> Xₐ
    γ, Xₐ --> 0
    β, Xₐ --> Xᵢ₁ + Xᵢ₂
    γ, Xᵢ₁ --> 0
    γ, Xᵢ₂ --> 0
 end C γ β
 jumpsys = convert(JumpSystem, rn, combinatoric_ratelaws = false)

C, γ, β, τ = [2., 0.1, 0.5, 15.]
p=[C, γ, β]
u0 = [0, 0, 0]
tf = 30.
saveat = .1
de_chan0 = [[],[]]
tspan = (0.,tf)
dprob = DiscreteProblem(u0, tspan,p)

# delay_trigger_affect! = function (integrator, rng)
#    append!(integrator.de_chan[1], τ)
#    append!(integrator.de_chan[2], τ)
# end
# delay_trigger = Dict(3=>delay_trigger_affect!)
delay_trigger = Dict(3=>[1=>τ,2=>τ])
delay_complete = Dict(1=>[2=>-1],2=>[3=>-1])

delay_affect1! = function (integrator, rng)
    i = rand(rng, 1:length(integrator.de_chan[1]))
    deleteat!(integrator.de_chan[1],i)
end
delay_affect2! = function (integrator, rng)
    i = rand(rng, 1:length(integrator.de_chan[2]))
    deleteat!(integrator.de_chan[2],i)
end
delay_interrupt = Dict(4=>delay_affect1!,5=>delay_affect2!)
delayjumpset = DelayJumpSet(delay_trigger,delay_complete,delay_interrupt)

djprob = DelayJumpProblem(jumpsys,dprob, DelayRejection(),  delayjumpset, de_chan0, save_positions=(false,false))
sol =@time solve(djprob, SSAStepper(),seed=10, saveat =.1)
plot(sol, label = [L"X_A" L"X_{I1}" L"X_{I2}"],legend =:topleft,xlabel="time",ylabel="Molecular Number")

timestamps = 0:0.1:tf
a = β + γ
x_A(t) = C/a*(1-exp(-a*t))
x_I(t)= 0<=t<=τ ? C*β/(a-γ)*((1-exp(-γ*t))/γ - (1-exp(-a*t))/a) : C*β/a*((1-exp(-γ*τ))/γ + exp(-a*t)*(1-exp((a-γ)τ))/(a-γ))
using Plots
plot(timestamps,x_A.(timestamps),linewidth=3,label=L"X_A",xlabel="time",ylabel="Mean Value")
plot!(timestamps,x_I.(timestamps),linewidth=3,label=L"X_I")

using StatsBase, DifferentialEquations.EnsembleAnalysis
Sample_size = Int(1e4)
ens_prob = EnsembleProblem(djprob)
ens =@time solve(ens_prob, SSAStepper(),EnsembleThreads(),trajectories = Sample_size, saveat = .1)

tsm(t) = timepoint_mean(ens,t)

mean_A(t) = tsm(t)[1]
mean_I1(t) = tsm(t)[2]
mean_I2(t) = tsm(t)[3]

plot(timestamps,x_A.(timestamps),linewidth=3,label=L"$X_A$ Exact",xlabel="time",ylabel="Mean Value")
plot!(timestamps,mean_A.(timestamps),linewidth=3,line=:dash,label=L"$X_A$ Delay SSA")
plot!(timestamps,x_I.(timestamps),linewidth=3,label=L"$X_I$ Exact")
plot!(timestamps,mean_I1.(timestamps),linewidth=3,line=:dash,label=L"$X_{I1}$ Delay SSA")
plot!(timestamps,mean_I2.(timestamps),linewidth=3,line=:dash, legend =:topleft,label=L"$X_{I2}$ Delay SSA")


using LaserLab
using Test
using Unitful
import Unitful:
    nm, μm, mm, cm, m, km,
    mg, g, kg,
    ps, ns, μs, ms, s, minute, hr, d, yr, Hz, kHz, MHz, GHz,
    eV,
    μJ, mJ, J,
	μW, mW, W

function hnu_eV(lambda_nm)
    return 1240.0 * (1.0/lambda_nm)
end

dspot(λ::Real, NA::Real) = 1.83*λ/(2*NA)

include("dffunctions_test.jl")
include("setup_test.jl")
include("glaser_test.jl")


# test data generation

using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
using IntervalTrees,SparseArrays,NearestNeighbors
using ViewerGL
GL = ViewerGL

function input_collection(data::Array)::Lar.LAR
	assembly = Lar.Struct(data)
	return Lar.struct2lar(assembly)
end

V,FV = Lar.sphere(2)([9,12])
#V,FV = Lar.sphere(2)([12,16])
EV = Lar.simplexFacets(FV)
#EV = Lar.quads2triangles(FV)
mysphere = V,FV,EV

data3d1 = mysphere
data3d2 = Lar.Struct([ Lar.t(0,1,0), mysphere ])
data3d3 = Lar.Struct([ Lar.t(0,0.5,0), Lar.s(0.4,0.4,0.4), mysphere ])
data3d4 = Lar.Struct([ Lar.t(4,0,0), Lar.s(0.8,0.8,0.8), mysphere ])
data3d5 = Lar.Struct([ Lar.t(4,0,0), Lar.s(0.4,0.4,0.4), mysphere ])

model3d = input_collection([ data3d5 ])#, data3d2, data3d3, data3d4, data3d5 ])
V,FV,EV = model3d
VV = [[k] for k in 1:size(V,2)];

GL.VIEW( GL.numbering(0.125)((V,[VV, EV, FV]),GL.COLORS[1],0.25) );

# copEV = Lar.coboundary_0(EV)
# copFE = Lar.coboundary_1(V,FV,EV)
# W = convert(Lar.Points, V')


using LorentzVectors
using Test
using Random

Random.seed!(8372946187652352328)

@testset "All tests" begin
    include("basics.jl")
    include("algebra.jl")
    include("random.jl")
    include("geometry.jl")
end;


"""
    break_stick_ibp(ν)

Break stick (IBP).
"""
function break_stick_ibp(ν)
    Kmax = size(ν, 1)
    p = []
    push!(p, ν[1,:])
    for k = 2:Kmax
        push!(p, p[k-1] .* ν[k,:])
    end
    p = hcat(p...)'
    return p
end

"""
    break_logstick_ibp(logν)

Break stick (IBP) in logarithmic space.
"""
function break_logstick_ibp(logν)
    Kmax = size(logν, 1)
    logp = []
    push!(logp, logν[1,:])
    for k = 2:Kmax
        push!(logp, logp[k-1] .+ logν[k,:])
    end
    logp = hcat(logp...)'
    return logp
end


function CliqueExpansion(H::hypergraph,weighted::Bool=true,binary::Bool=false)
    """
    Weighted clique expansion where a hyperedge e is expanded to a
    weighted clique with each edge having weight 1/(|e| - 1)
    """
    n = length(H.D)
    I = Int64[]
    J = Int64[]
    V = Float64[]
    ks = setdiff(keys(H.E),1)
    for k in ks
        for edge in keys(H.E[k])
            weight = H.E[k][edge]
            for i = 1:k-1
                ei = edge[i]
                for j = i+1:k
                    ej = edge[j]
                    push!(I,ei)
                    push!(J,ej)
                    if weighted
                        push!(V, weight / (k - 1))
                    else
                        push!(V, weight)
                    end
                end
            end
        end
    end
    A = SparseArrays.sparse(I,J,V,n,n)
    for i = 1:n
        A[i, i] = 0.0
    end
    SparseArrays.dropzeros!(A)
    A = SparseArrays.sparse(A+A')
    if binary
        I, J, V = SparseArrays.findnz(A)
        A = SparseArrays.sparse(I, J, 1, n, n)
    end
    return A
end

function CliqueExpansionModularity(H::hypergraph,gamma::Float64=1.0,weighted::Bool=true,randflag::Bool=false,binary::Bool=false,clusterpenalty::Float64=0.0,maxits::Int64=10000)
    return CliqueExpansionModularity(H,gamma;weighted=weighted,randflag=randflag,binary=binary,clusterpenalty=clusterpenalty,maxits=maxits)
end

function CliqueExpansionModularity(H::hypergraph,gamma::Float64=1.0;weighted::Bool=true,randflag::Bool=false,binary::Bool=false,clusterpenalty::Float64=0.0,maxits::Int64=10000)
    """
    Perform a clique expansion on the hypergraph H and then run vanilla
    modularity on the resulting graph.
    """
    A = CliqueExpansion(H,weighted,binary)
    return VanillaModularity(A,gamma,randflag,clusterpenalty,maxits)
end

function StarExpansionModularity(H::hypergraph,gamma::Float64=1.0;weighted::Bool=true,randflag::Bool=false,binary::Bool=false,maxits::Int64=100)
    """
    Perform a clique expansion on the hypergraph H and then run vanilla
    modularity on the resulting graph.
    """
    He2n, w = hypergraph2incidence(H)
    m,n = size(He2n)
    A = [spzeros(n,n) He2n'; He2n spzeros(m,m) ]
    Za = VanillaModularity(A,gamma,randflag,maxits)
    return Za[1:n]
end

function VanillaModularity(A::SparseArrays.SparseMatrixCSC{Float64,Int64},gamma::Float64=1.0,randflag::Bool=false,clusterpenalty::Float64=0.0,maxits::Int64=10000)
    """
    Vanilla modularity algorithm, obtained by calling the LambdaLouvain algorithm
    implementation from:

    Parameterized Correlation Clustering in Hypergraphs and Bipartite Graphs
    https://arxiv.org/abs/2002.09460

    Code: https://github.com/nveldt/ParamCC/blob/master/src/Graph_Louvain.jl
    """

    d = vec(sum(A,dims = 2))
    n = length(d)
    vol = sum(d)
    lam = gamma/vol
    Cs = LambdaLouvain(A,d,lam,randflag,maxits,clusterpenalty)

    c = Cs[:,end]
    @assert(length(c) == n)

    return c
end

function computeDyadicResolutionParameter(H, Z; mode = "γ", weighted=true, binary=false, pseudocount = 0.0)
    """
    compute the dyadic resolution parameter associated to a partition using the formula from Newman (2016): https://arxiv.org/abs/1606.02319
    """

    G = CliqueExpansion(H, weighted, binary)
    I, J = SparseArrays.findnz(G)

    n = length(H.D)  # number of nodes
    m = sum(G)/2     # number of edges

    # form degree sequence and edge counts
    D = vec(sum(G, dims=1))

    m_in = pseudocount
    m_out = pseudocount

    for k in 1:length(I)
        if Z[I[k]] == Z[J[k]]
            m_in  += G[I[k], J[k]]/2
        else
            m_out += G[I[k], J[k]]/2
        end
    end

    # compute resolution parameter
    V = [sum(D[findall(Z .== c)]) for c in unique(Z)]
    ωᵢ = 4*m*m_in / (sum(V.^2))
    ωₒ = (2m_out)/(2m - (sum(V.^2)/(2m)))

    if mode == "γ"
        γ = (ωᵢ - ωₒ)/(log(ωᵢ) - log(ωₒ))
        return γ
    else
        return(ωᵢ, ωₒ)
    end
end

function dyadicModularity(H, Z, γ; weighted=false, binary=false)
    """
    Compute the normalized dyadic modularity a given partition in the graph obtained from projecting a hypergraph.
    The resolution parameter and basic properties of the dyadic graph projection can be specified.

    # Arguments

    - H::hypergraph, a hypergraph
    - Z::Vector{Int4}, a label vector
    - γ::Float64, the resolution parameter. Must be strictly positive.
    - weighted::bool, whether to weight the clique obtained by projecting a hypergraph according to a factor of 1/(k-1), where k is the clique size. Choosing to do so preserves the degrees of nodes in the dyadic projection. If false and if binary is false, then each edge within the clique has weight 1.
    - binary::bool, whether to convert the graph to a binary graph in which present edges have weight 1 and absent edges have weight 0. If false and if weighted is false, then the weight of an edge between two nodes in the return value is equal to the number of hyperedges containing both of those nodes.

    # Returns

    - Q::Float64, a normalized modularity score between -1 and 1.

    """
    G = CliqueExpansion(H, weighted, binary)
    d = vec(sum(G, dims=1))

    # non-degree (cut) term
    edge_obj = 0.0
    for (i, j, v) in zip(SparseArrays.findnz(G)...)
        if Z[i] == Z[j]
            edge_obj += v
        end
    end

    # volume terms
    vols = Dict{Int64, Float64}()
    for c in unique(Z)
        vols[c] = 0.0
    end
    for i = 1:length(d)
        vols[Z[i]] += d[i]
    end

    Q = edge_obj
    volG = sum(d)
    vol_term = 0.0
    for c in unique(Z)
        Q -= γ * vols[c]^2 / volG
    end

    return Q / volG
end

function dyadicLogLikelihood(H, Z, ωᵢ, ωₒ; weighted=false, binary=false, constants = false)
    G = CliqueExpansion(H, weighted, binary)
    d = vec(sum(G, dims=1))

    # Eq. (14) from https://arxiv.org/pdf/1606.02319.pdf
    γ = (ωᵢ - ωₒ) / (log(ωᵢ) - log(ωₒ))
    Q = dyadicModularity(H, Z, γ; weighted=weighted)
    m = sum(d) / 2
    B = m * log(ωᵢ / ωₒ)
    C = m * (ωₒ + log(ωₒ))

    if constants
        K = sum(SpecialFunctions.loggamma.(SparseArrays.nonzeros(G).+1)) / 2
        return B * Q + C - K
    end
    return B * Q + C
end

function dyadicLogLikelihoodDirect(H, Z, ωᵢ, ωₒ)
    """
    A naive version of the dyadic log-likelihood calculation, used only for testing the fast one.
    """
    G = CliqueExpansion(H, false, false)
    d = vec(sum(G, dims=1))
    m = sum(d)/2
    L = 0.0
    for (i, j, v) in zip(SparseArrays.findnz(G)...)
        if i < j
            a = G[i,j]
            ω = Z[i] == Z[j] ? ωᵢ : ωₒ
            L += log(poisson_pdf(Int(a), ω*d[i]*d[j]/(2m)))
        end
    end
    return L
end

# Below are functions for just graph louvain, nothing about hypergraphs

function ConstructAdj(C::SparseArrays.SparseMatrixCSC,n::Int64)
    """
    ConstructAdj: Construct Adjacency List
    This takes in a sparse adjacency matrix for a graph, and returns an adjacency
    list. While it seems inefficient to store the matrix in multiple ways, as long
    as there are no terrible memory issues, it is VERY handy to be able to quickly
    get a list of neighbors for each node.

    The function also returns the degree vector d. This is NOT the weighted degree,
        d[i] = total number of neighbors of node i
    """
    rp = C.rowval
    ci = C.colptr
    Neighbs = Vector{Vector{Int64}}()
    d = zeros(Int64,n)
    for i = 1:n
        # chop up the rp vector and put it in Neighbs
        push!(Neighbs,rp[ci[i]:ci[i+1]-1])
        d[i] = ci[i+1]-ci[i]
    end

    # d is the number of neighbors. This is the unweighted degree,
    # but note importantly that if the original graph is weighted this is
    # not the same as the degree vector d we will sometimes use
    return Neighbs, d
end

"""
Evaluates the LambdaCC objective, which is equivalent to modularity in graphs.

https://dl.acm.org/doi/pdf/10.1145/3178876.3186110

A = adjacency matrix for a graph
c = clustering vector
w = weights vector (often taken to be degrees of nodes)
lam = resolution paramter

This can be replaced with our function for checking the hypergraph modularity
objective.
"""
function LamCCobj(A::SparseArrays.SparseMatrixCSC,c,w,lam)

    w_volA = sum(w) # weighted volume
    obj = (lam*(w_volA)^2-lam*sum(w.^2))/2 # constant with respect to the clustering

    for i = 1:maximum(c)
        S = findall(x->x==i,c)
        AS = A[:,S]
        vol = sum(AS.nzval);
        SAS = AS[S,:]
        edges = sum(SAS.nzval);
        cut = vol-edges
        w_volS = sum(w[S])
        obj += 1/2*(cut - lam*w_volS*(w_volA-w_volS))
    end
    return obj
end

"""
Run the Louvain algorithm many times, taking the result with the best objective.
"""
function Many_Louvain(A::SparseArrays.SparseMatrixCSC{Float64,Int64},w::Vector{Float64},lam::Float64,numtimes::Int64,maxits::Int64=10000)

    n = size(A,1)
    BestObj = Inf
    cBest = collect(1:n)

    for k = 1:numtimes

        ######
        Cs = LambdaLouvain(A,w,lam,true,maxits)
        ######

        c = Cs[:,end]

        ######
        obj = LamCCobj(A,c,w,lam)
        ######

        if obj < BestObj
            BestObj = obj
            cBest = c
        end
    end

    return cBest, BestObj
end

"""
The full Louvain algorithm. Lambda is the resolution parameter; think of it as
lambda = gamma/(vol(G)), where gamma is the more traditional resolution parameter
associated with the modularity objective.

w = weights function, w[i] = weight for node i. This is often the degree of node
    i, but doesn't have to be.
"""
function LambdaLouvain(A::SparseArrays.SparseMatrixCSC{Float64,Int64},w::Vector{Float64},lam::Float64,randflag::Bool=false,maxits::Int64=10000,clusterpenalty::Float64=0.0)

    @assert(LinearAlgebra.issymmetric(A))
    n = size(A,1)

    # Step 1: greedy moves until no more improvement
    c, improved = LambdaLouvain_Step(A,w,lam,randflag,maxits,clusterpenalty)

    # Keeps track of the clustering that is found after each call to Step 1.
    if improved
        Cs = c
        c_old = copy(c)
    else
        Cs = c
    end

    # As long as something is still improving each time you call step 1, keep going
    while improved

        # Step 2: Collapse the clustering into supernodes

        Anew, wnew = collapse_clustering(A,w,c_old)

        # Step 1: Go back to greedy local moves, this time on the reduced graph
        cSuper, improved = LambdaLouvain_Step(Anew,wnew,lam,randflag,maxits,clusterpenalty)
        N = length(wnew)    # N = number of supernodes = number clusters from last round

        # Undo the procedure that merged nodes into super nodes, so that you get a clustering vector of length n
        # Do this only if the last call to Step 1 led to at least one greedy move.

        if improved

            # Extract what that new clustering is
            c_new = zeros(Int64,n)

            Clusters = clusters_from_cvec(c_old)

            # For each supernode, place all original node IDs that make it up
            # into a cluster of the supernodes label
            for i = 1:N

                # Get the cluster that supernode i is in
                SuperI_Cluster = cSuper[i]

                # Get individual node ID that are in supernode i.

                SuperI_nodes = Clusters[i]
                c_new[SuperI_nodes] .= SuperI_Cluster
            end
            Cs = [Cs c_new]
            c_old = copy(c_new)
        end
    end

    return Cs
end

"""
From a cluster indicator vector, extract a vector of vectors storing clusters
"""

function clusters_from_cvec(c::Vector{Int64})

    Clusters = Vector{Vector{Int64}}()
    for v = 1:maximum(c)
        push!(Clusters, Vector{Int64}())
    end
    for i = 1:length(c)
        push!(Clusters[c[i]],i)
    end
    return Clusters
end

"""
Run Step 1 of the Louvain algorithm: iterate through nodes and greedily move
nodes to adjacent clusters.
"""
function LambdaLouvain_Step(A::SparseArrays.SparseMatrixCSC{Float64,Int64},w::Vector{Float64},lam::Float64,randflag::Bool=false,maxits::Int64=Inf,clusterpenalty::Float64=0.0)
    n = size(A,1)

    # This permutes the node labels, to add randomization in the Louvain
    # algorithm so that you don't always traverse the nodes in the same order
    if randflag
        p = Random.randperm(n)
        A = A[p,p]
        undop = sortperm(p)
        w = w[p]
    end

    # All nodes start in singleton clusters
    c = collect(1:n)
    @assert(size(w,1) == n)
    improving = true
    its = 0

    # number of clusters
    K = n

    # Neighs[i] lists nodes that share an edge with node i
    # degvec[i] = number of neighbors of node i
    Neighbs, degvec = ConstructAdj(A,n)

    # Store a vector of n clusters
    Clusters = Vector{Vector{Int64}}()
    for v = 1:n
        push!(Clusters, Vector{Int64}())
        push!(Clusters[v],v)    # singleton clusters--one for each node
    end

    # As long as we can keep improving, continue
    while improving && its < maxits
        its += 1
        # println("Iteration $its")
        improving = false
        nextclus = 1

        # visit each node in turn
        for i = 1:n

            # Cluster index for node i
            Ci_ind = c[i]

            # Get the indices of nodes in i's cluster
            Ci = Clusters[Ci_ind]

            clus_size = length(Ci)

            Ni = Neighbs[i]

            ## First few lines here essentially compute the current contribution
            # this node i has on the objective, based on its current cluster placement

            # Weight of negative mistakes currently at i
            neg_inner = w[i]*(sum(w[Ci]) - w[i])

            # Weight of positive mistakes if we remove i from Ci
            pos_inner = sum(A[i,Ci])

            # Increase in mistakes if we remove i from Ci
            total_inner = pos_inner - lam*neg_inner

            BestImprove = 0
            BestC = Ci_ind

            Ai = A[:,i]

            # Get the neighboring clusters of i:
            # there are the only places we would even consider moving node i to
            NC = unique(c[Ni])

            # Now let's see if it's better to move to a nearby cluster, Cj
            for j = 1:length(NC)
                Cj_ind = NC[j]

                # Check how much it would improve to move i to to cluster j
                if Cj_ind == Ci_ind
                    change = 0
                else

                    Cj = Clusters[Cj_ind]

                    # Moving i from Ci to Cj adds negative mistakes
                    neg_outer = w[i]*(sum(w[Cj]))

                    # Moving i from Ci to Cj decreases positive mistakes

                    pos_outer = mysum(Ai,Cj)

                    total_outer = lam*neg_outer - pos_outer

                    # Change in cluster part of objective
                    if clus_size == 1
                        Δclus = clusterpenalty*(log(K-1)-log(K))
                    else
                        Δclus = 0
                    end

                    # This is the overall change in objective if we move i to Cj
                    change = total_outer + total_inner + Δclus

                end

                # Check if this is currently the best possible greedy move to make
                if change < BestImprove
                    BestImprove = change
                    BestC = Cj_ind
                    improving = true
                end
            end
            # Done checking whether there's a better place to greedily move the cluster to.

            # Move i to the best cluster to move it to, do this
            # only if it strictly improves the objective
            if BestImprove < 0
                ci_old = c[i]
                c[i] = BestC

                # Remove i from its old cluster...
                Clusters[ci_old] = setdiff(Clusters[ci_old],i)

                # ...and add it to its new cluster
                push!(Clusters[BestC],i)
                sort!(Clusters[BestC])

                improving = true # we have a reason to keep iterating!
                if clus_size == 1
                    # we moved a singleton to another cluster
                    K -= 1
                end
            end

        end
    end
    if its == 1
        improved = false
    else
        improved = true
        c, Clusters = renumber(c,Clusters)
    end

    if randflag
        c = c[undop]    # if you previously mixed up node order, re-order them now
    end

    return c, improved

end

function mysum(v::SparseArrays.SparseVector{Float64,Int64},inds::Vector{Int64})
    """
    Less numerically stable, but for the purposes of a greedy clustering
    algorithm works as well as it needs to, and is faster.
    """
    s = 0
    for i in inds
        s += v[i]
    end
    return s
end

"""
Step 2 of the Louvain algorithm:
Collapse a clustering into a new network of supernodes and weighted edges, so
that you can then run the same greedy node-moving on supernodes.
"""
function collapse_clustering(A::SparseArrays.SparseMatrixCSC{Float64,Int64},w::Vector{Float64},c::Vector{Int64})

    n = size(A,1)
    Clusters = Vector{Vector{Int64}}()
    ClusterNeighbs = Vector{Set{Int64}}()
    Neighbs, degvec = ConstructAdj(A,n)

    for v = 1:maximum(c)
        push!(Clusters, Vector{Int64}())
        push!(ClusterNeighbs, Set{Int64}())
    end

    for i = 1:n
        push!(Clusters[c[i]],i)
        # Also keep track of which clusters are adjacent clusters
        for j in Neighbs[i]
            if c[j] != c[i]
                push!(ClusterNeighbs[c[i]],c[j])
            end
        end
    end

    # Number of supernodes to form = number of clusters
    N = round(Int64,maximum(c))

    I = Vector{Int64}()
    J = Vector{Int64}()
    V = Vector{Float64}()
    wnew = zeros(N)
    # Construct a new sparse matrix with new node weights
    for i = 1:N
        Ci = sort(Clusters[i])
        wnew[i] = sum(w[Ci])
        ACi = A[:,Ci]
        for j in ClusterNeighbs[i]
            Cj = Clusters[j]
            Eij = sum(ACi[Cj,:])
            push!(I,i)
            push!(J,j)
            push!(V,Eij)
        end
    end

    Anew = SparseArrays.sparse(I,J,V,N,N)
    Anew = Anew+Anew'

    return Anew, wnew
end

# Sort sizes: Return a set of cluster sizes, arranged in order
function ClusterSizes(c)
    c = renumber(c)

    sizes = Vector{Int64}()
    for i = 1:maximum(c)
        inds = findall(x->x==i,c)
        push!(sizes,length(inds))
    end

    return sort(sizes,rev=true)
end


#sbox5.jl
"""Implement a pattern for dispatching actions based on Truth table values"""

module DispatchPattern
using Match

struct Truth
  a::Bool
  b::Bool
end

actionAB()   =println("Do A and B")
actionA!B()  =println("Do A not B")
action!AB()  =println("Do B not A")
action!A!B() =println("Do nothing")

"Pattern for dispatching based on multiple truth cases"
dispatch!(item::Truth)=@match (item.a, item.b)  begin
  (true, true)          => actionAB()
  (true, false)         => actionA!B()
  (false, true)         => action!AB()
  (false, false)        => action!A!B()
  bad                   => println("Unknown dispatch: $bad")
end

truths=[Truth(true, true),
        Truth(true, false),
        Truth(false, true),
        Truth(false, false)]

map(dispatch!, truths)

end


module ContinuousTransformations

import Base.rand

using ArgCheck
using AutoHashEquals
import Distributions:
    ContinuousUnivariateDistribution, ContinuousMultivariateDistribution, logpdf
using DocStringExtensions
using LinearAlgebra
using MacroTools
using Parameters
using StatsFuns
using Lazy
# using Unrolled

export
    # intervals
    AbstractInterval, RealLine, ℝ, PositiveRay, ℝ⁺, NegativeRay, ℝ⁻,
    Segment, width,
    # general API
    domain, image, isincreasing, transform, inverse, logjac,
    transform_and_logjac, @increment_logjac, ContinuousTransformation,
    # univariate transformations
    UnivariateTransformation,
    Affine, IDENTITY,
    Negation, NEGATION, Logistic, LOGISTIC, RealCircle, REALCIRCLE, Exp, EXP,
    Logit, LOGIT, InvRealCircle, INVREALCIRCLE, Log, LOG,
    affine_bridge, default_transformation, bridge,
    # multivariate transformations NOTE API may change
    UnitVector, CorrelationCholeskyFactor, lkj_correlation_cholesky_logpdf,
    # grouped transformations
    GroupedTransformation, get_transformation, ArrayTransformation,
    TransformationTuple,
    # wrapped transformations
    TransformationWrapper, TransformLogLikelihood, get_loglikelihood,
    TransformDistribution, get_distribution, logpdf_in_domain, logpdf_in_image,
    # utilities
    ungrouping_map, all_finite

import Base: in, length, size, ∘, show, getindex, intersect, extrema, minimum,
    maximum, isfinite, isinf, isapprox

import Statistics: middle

# utilities

"""
    $SIGNATURES

Define an `isapprox` method, comparing the given fields in type `T`.
"""
macro define_isapprox(T, fields...)
    body = foldl((a, b) -> :($(a) && $(b)),
                 :(isapprox(x.$(field), y.$(field); rtol=rtol, atol=atol))
                 for field in fields)
    quote
        function Base.isapprox(x::$T, y::$T; rtol::Real=√eps(), atol::Real=0)
            $(body)
        end
    end
end

"""
    $SIGNATURES

Define a singleton type with the given name and supertype (specified
as `name <: supertype`), and a constant which defaults to the name in
uppercase.
"""
macro define_singleton(name_and_supertype, constant = nothing)
    @capture name_and_supertype name_ <: supertype_
    if constant == nothing
        constant = Symbol(uppercase(string(name)))
    end
    quote
        Core.@__doc__ struct $(esc(name)) <: $(esc(supertype)) end
        Core.@__doc__ const $(esc(constant)) = $(esc(name))()
    end
end

"""
    _fma(x, y, z)

Placeholder for `Base.fma` until
https://github.com/JuliaDiff/ReverseDiff.jl/issues/86 is fixed.
"""
_fma(x, y, z) = x*y+z

"""
    $SIGNATURES

Define [`transform`](@ref) and [`logjac`](@ref) using
[`transform_and_logjac`](@ref).
"""
macro define_from_transform_and_logjac(T)
    S = esc(T)
    quote
        (t::$S)(x) = transform_and_logjac(t, x)[1]
        ContinuousTransformations.logjac(t::$S, x) = transform_and_logjac(t, x)[2]
    end
end

# abstract interface for transformations

"""
    logjac(t, x)

The log of the determinant of the Jacobian of `t` at `x`.
```
"""
function logjac end

"""
    inverse(t, x)

Return ``t⁻¹(x)``.

    inverse(t)

Return the transformation ``t⁻¹``.
"""
inverse(t, x) = inverse(t)(x)

"""
    domain(transformation)

Return the domain of the transformation.
"""
function domain end

"""
    image(transformation)

Return the image of the transformation.
"""
function image end

"""
$TYPEDEF

Continuous bijection ``D ⊂ ℝ^n→ I ⊂ ℝ^n`` or ``D ⊂ ℝ → I ⊂ ℝ``.
"""
abstract type ContinuousTransformation <: Function end

"""
    transform(t, x)

Return ``t(x)``. NOTE: this equivalent to the callable syntax `t(x)`, which is
also available.
"""
transform(t::ContinuousTransformation, x) = t(x)

"""
    y, logjac = $SIGNATURES

Return the transformed value and the determinant of the log Jacobian.

Equivalent to `t(x), logjac(t, x)`, but may be faster because of reused
components.
"""
function transform_and_logjac(t::ContinuousTransformation, x)
    transform(t, x), logjac(t, x)
end

"""
    $SIGNATURES

Return `transformation(x)` while incrementing `logjac` with the log Jacobian
determinant.
"""
macro increment_logjac(logjac, transformation, x)
    quote
        y, lj = transform_and_logjac($(esc(transformation)), $(esc(x)))
        $(esc(logjac)) += lj
        y
    end
end

"""
    $SIGNATURES

Return `transformation(x)` while incrementing `logjac` with the log Jacobian
determinant.
"""
macro increment_logjac(logjac, transformation_x)
    @capture transformation_x transformation_(x_)
    quote
        @increment_logjac $(esc(logjac)) $(esc(transformation)) $(esc(x))
    end
end

"""
    rhs_string(transformation, term)

Return the formula representing the hand side of the `transformation`, with
`term` as the argument.
"""
function rhs_string end

transformation_string(t, x = "x") = x * " ↦ " * rhs_string(t, x)

# this is necessary because `display` calls this method, which is defined for
# Function and would match
show(io::IO, ::MIME"text/plain", t::ContinuousTransformation) =
    println(io, transformation_string(t, "x"))

show(io::IO, t::ContinuousTransformation) =
    print(io, transformation_string(t, "x"))

"""
    bridge(dom, img, [transformation])

Return a transformation that maps `dom` to `img`.

The `transformation` argument may be used to specify a particular transformation
family, otherwise `default_transformation` is used.
"""
function bridge end

# intervals

"""
$TYPEDEF

Abstract supertype for all univariate intervals. It is not specified whether
they are open or closed.
"""
abstract type AbstractInterval end

isinf(x::AbstractInterval) = !isfinite(x)

isapprox(::AbstractInterval, ::AbstractInterval; rtol=√eps(), atol=0) = false

extrema(x::AbstractInterval) = minimum(x), maximum(x)

"""
    RealLine()

The real line. Use the constant [`ℝ`](@ref).
"""
struct RealLine <: AbstractInterval end

isapprox(::RealLine, ::RealLine; rtol=√eps(), atol=0) = true

show(io::IO, ::RealLine) = print(io, "ℝ")

"A constant for the real line."
const ℝ = RealLine()

in(x::Real, ::RealLine) = true

minimum(::RealLine) = -Inf
maximum(::RealLine) = Inf

isfinite(::RealLine) = false

"""
    PositiveRay(left)

The real numbers above `left`. See [`ℝ⁺`](@ref).
"""
@auto_hash_equals struct PositiveRay{T <: Real} <: AbstractInterval
    left::T
    function PositiveRay{T}(left::T) where T
        @argcheck isfinite(left) "Need finite endpoint."
        new(left)
    end
end

PositiveRay(left::T) where {T} = PositiveRay{T}(left)

in(x::Real, ray::PositiveRay) = ray.left ≤ x
minimum(ray::PositiveRay) = ray.left
maximum(::PositiveRay{T}) where T = T(Inf)
isfinite(::PositiveRay) = false

"The positive real numbers."
const ℝ⁺ = PositiveRay(0.0)

@define_isapprox PositiveRay left

"""
    NegativeRay(right)

The real numbers below `right`. See [`ℝ⁻`](@ref).
"""
@auto_hash_equals struct NegativeRay{T <: Real} <: AbstractInterval
    right::T
    function NegativeRay{T}(right::T) where T
        @argcheck isfinite(right) "Need finite endpoint."
        new(right)
    end
end

NegativeRay(right::T) where {T} = NegativeRay{T}(right)

in(x::Real, ray::NegativeRay) = x ≤ ray.right
minimum(::NegativeRay{T}) where T = -T(Inf)
maximum(ray::NegativeRay) = ray.right
isfinite(::NegativeRay) = false

"The negative real numbers."
const ℝ⁻ = NegativeRay(0.0)

@define_isapprox NegativeRay right

"""
    Segment(left, right)

The real numbers between `left` and `right`, with
``-∞ < \\text{left} < \\text{right} < ∞`` enforced.
"""
@auto_hash_equals struct Segment{T <: Real} <: AbstractInterval
    left::T
    right::T
    function Segment{T}(left::T, right::T) where T
        @argcheck isfinite(left) && isfinite(right) "Need finite endpoints."
        @argcheck left < right "Need strictly increasing endpoints."
        new(left, right)
    end
end

Segment(left::T, right::T) where {T <: Real} = Segment{T}(left, right)

Segment(left::Real, right::Real) = Segment(promote(left, right)...)

in(x::Real, s::Segment) = s.left ≤ x ≤ s.right
minimum(s::Segment) = s.left
maximum(s::Segment) = s.right
isfinite(::Segment) = true

@define_isapprox Segment left right

"""
    $SIGNATURES

Width of a finite interval.
"""
width(s::Segment) = s.right - s.left

middle(s::Segment) = middle(s.left, s.right)

# intersections

## general fallback method. define specific methods with the following
## argument ordering Segment < PositiveRay < NegativeRay < RealLine < all
intersect(a::AbstractInterval, b::AbstractInterval) = intersect(b, a)

intersect(a::RealLine, b::AbstractInterval) = b

"""
    $SIGNATURES

Helper function for forming a segment when possible. Internal, not exported.
"""
@inline function _maybe_segment(a, b)
    # NOTE Decided not to represent the empty interval, as it has no use in
    # the context of this package. Best to throw an error as soon as possible.
    a < b ? Segment(a, b) :
        throw(ArgumentError("intersection of intervals is empty"))
end

intersect(a::Segment, b::Segment) =
    _maybe_segment(max(a.left, b.left), min(a.right, b.right))

intersect(a::Segment, b::PositiveRay) =
    _maybe_segment(max(b.left, a.left), a.right)

intersect(a::Segment, b::NegativeRay) =
    _maybe_segment(a.left, min(a.right, b.right))

intersect(a::PositiveRay, b::PositiveRay) = PositiveRay(max(a.left, b.left))

intersect(a::PositiveRay, b::NegativeRay) = _maybe_segment(a.left, b.right)

intersect(a::NegativeRay, b::NegativeRay) = NegativeRay(min(a.right, b.right))

# univariate transformations

"""
$TYPEDEF

Univariate monotone transformation, either *increasing* or *decreasing* on the
whole domain (thus, a bijection).
"""
abstract type UnivariateTransformation <: ContinuousTransformation end

length(::UnivariateTransformation) = 1

size(::UnivariateTransformation) = ()

"""
    isincreasing(transformation)

Return `true` (`false`), when the transformation is monotonically increasing
(decreasing).
"""
function isincreasing end

"""
$TYPEDEF

Trait that is useful for domain and image calculations. See [`RRStable`](@ref).
"""
abstract type RRStability end

"""
$TYPEDEF

Trait that indicates that a univariate transformation

1. maps ``ℝ`` to ``ℝ``,

2. supports mapping intervals, and

3. maps subtypes of [`AbstractInterval`](@ref) to the same type.
"""
struct RRStable <: RRStability end

"""
$TYPEDEF

Trait that indicates that a univariate transformation is *not* [`RRStable`](@ref).
"""
struct NotRRStable <: RRStability end

"""
    $SIGNATURES

Return either the trait [`RRStable`](@ref) and `NotRRStable`.
"""
RR_stability(::UnivariateTransformation) = NotRRStable()

∘(::RRStability, ::RRStability) = NotRRStable()

∘(::RRStable, ::RRStable) = RRStable()

# Affine and Negation
#
# These play a special role since map subtypes of `AbstractInterval` to
# themselves, and thus are useful for composition.

"""
    Affine(α, β)

Mapping ``ℝ → ℝ`` using ``x ↦ α⋅x + β``.

``α > 0`` is enforced, see [`Negation`](@ref).
"""
@auto_hash_equals struct Affine{T <: Real} <: UnivariateTransformation
    α::T
    β::T
    function Affine{T}(α::T, β::T) where T
        @argcheck α > 0 DomainError
        new(α, β)
    end
end

Affine(α::T, β::T) where T = Affine{T}(α, β)
Affine(α, β) = Affine(promote(α, β)...)

"Identity (as an affine transformation)."
const IDENTITY = Affine(1, 0)

domain(::Affine) = ℝ
image(::Affine) = ℝ
(t::Affine)(x) = _fma(x, t.α, t.β)
logjac(t::Affine, x) = log(abs(t.α))
inverse(t::Affine) = Affine(1/t.α, -t.β/t.α)
isincreasing(t::Affine) = true

RR_stability(::Affine) = RRStable()

(t::Affine)(x::Segment) = Segment(t(x.left), t(x.right))
(t::Affine)(x::PositiveRay) = PositiveRay(t(x.left))
(t::Affine)(x::NegativeRay) = NegativeRay(t(x.right))
(t::Affine)(::RealLine) = ℝ

function rhs_string(t::Affine, term)
    @unpack α, β = t
    α == 1 || (term = "$(round(α; sigdigits = 4))⋅" * term)
    β == 0 || (term = term * " + $(round(β, sigdigits = 4))")
    term
end

"""
    Negation()

Mapping ``ℝ → ℝ`` using ``x ↦ -x``.
"""
@define_singleton Negation <: UnivariateTransformation

domain(::Negation) = ℝ
image(::Negation) = ℝ
(::Negation)(x) = -x
logjac(::Negation, x) = zero(x)
inverse(::Negation) = NEGATION
isincreasing(::Negation) = false

RR_stability(::Negation) = RRStable()

(::Negation)(x::Segment) = Segment(-x.right, -x.left)
(::Negation)(x::PositiveRay) = NegativeRay(-x.left)
(::Negation)(x::NegativeRay) = PositiveRay(-x.right)
(::Negation)(::RealLine) = ℝ

rhs_string(::Negation, term) = "-" * term

# transformations from ℝ to subsets

"""
    Logistic()

Mapping ``ℝ → (0,1)`` using ``x ↦ 1/(1+\\exp(-x))``.
"""
@define_singleton Logistic <: UnivariateTransformation

domain(::Logistic) = ℝ
image(::Logistic) = Segment(0, 1)
(t::Logistic)(x) = logistic(x)
logjac(::Logistic, x) = -(log1pexp(x)+log1pexp(-x))
inverse(::Logistic) = LOGIT
isincreasing(::Logistic) = true

rhs_string(::Logistic, term) = "logistic($term)"

"""
    RealCircle()

Mapping ``ℝ → (-1,1)`` using ``x ↦ x/√(1+x^2)``.
"""
@define_singleton RealCircle <: UnivariateTransformation

domain(::RealCircle) = ℝ
image(::RealCircle) = Segment(-1, 1)
(t::RealCircle)(x) = isinf(x) ? sign(x) : x/√(1+x^2)
logjac(::RealCircle, x) = -1.5*log1psq(x)
inverse(::RealCircle) = INVREALCIRCLE # x/√(1-x^2)
isincreasing(::RealCircle) = true

rhs_string(::RealCircle, term) = "realcircle($term)"

"""
    Exp()

Mapping ``ℝ → ℝ⁺`` using ``x ↦ \\exp(x)``.
"""
@define_singleton Exp <: UnivariateTransformation

domain(::Exp) = ℝ
image(::Exp) = ℝ⁺
(t::Exp)(x) = exp(x)
logjac(::Exp, x) = x
inverse(::Exp) = LOG
isincreasing(::Exp) = true

rhs_string(::Exp, term) = "exp($term)"

# transformations to ℝ from subsets

"""
    Logit()

Mapping ``(0,1) → ℝ`` using ``x ↦ \\log(x/(1-x))``.
"""
@define_singleton Logit <: UnivariateTransformation

domain(::Logit) = Segment(0, 1)
image(::Logit) = ℝ
(t::Logit)(x) = logit(x)
logjac(::Logit, x) = -log(x)-log1p(-x)
inverse(::Logit) = LOGISTIC
isincreasing(::Logit) = true

rhs_string(::Logit, term) = "logit($term)"

"""
    InvRealCircle()

Mapping ``(-1,1) → ℝ`` using ``x ↦ x/√(1-x^2)``.
"""
@define_singleton InvRealCircle <: UnivariateTransformation

domain(::InvRealCircle) = Segment(-1, 1)
image(::InvRealCircle) = ℝ
(t::InvRealCircle)(x) = x/√(1-x^2)
logjac(::InvRealCircle, x) = -1.5*log1p(-abs2(x))
inverse(::InvRealCircle) = REALCIRCLE
isincreasing(::InvRealCircle) = true

rhs_string(::InvRealCircle, term) = "realcircle⁻¹($term)"

"""
    Log()

Mapping ``ℝ → ℝ⁺`` using ``x ↦ \\exp(x)``.
"""
@define_singleton Log <: UnivariateTransformation

domain(::Log) = ℝ⁺
image(::Log) = ℝ
(t::Log)(x) = log(x)
logjac(::Log, x) = -log(x)
inverse(::Log) = EXP
isincreasing(::Log) = true

rhs_string(::Log, term) = "log($term)"

# composed transformations

"""
    ComposedTransformation(f, g)

Compose two univariate transformations, resulting in the mapping
``f∘g``, or ``x ↦ f(g(x))`.

Use the `∘` operator for construction.
"""
struct ComposedTransformation{Tf <: UnivariateTransformation,
                              Tg <: UnivariateTransformation} <:
                                  UnivariateTransformation
    f::Tf
    g::Tg
end

RR_stability(t::ComposedTransformation) = RR_stability(t.f) ∘ RR_stability(t.g)

(c::ComposedTransformation)(x) = c.f(c.g(x))

rhs_string(t::ComposedTransformation, term) =
    rhs_string(t.f, rhs_string(t.g, term))

# composed image and domain uses traits, effectively allowing calculation for
# composing with Affine and Negation

"""
    composed_image(f_RR_stability, g_RR_stability, f, g)
"""
composed_image(::RRStable, ::RRStability, f, g) = f(image(g))

composed_image(::RRStability, ::RRStable, f, g) = image(f)

composed_image(::RRStable, ::RRStable, f, g) = ℝ

"""
    composed_domain(f_RR_stability, g_RR_stability, f, g)
"""
composed_domain(::RRStable, ::RRStability, f, g) = domain(g)

composed_domain(::RRStability, ::RRStable, f, g) = inverse(g)(domain(f))

composed_domain(::RRStable, ::RRStable, f, g) = ℝ

image(t::ComposedTransformation) =
    composed_image(RR_stability(t.f), RR_stability(t.g), t.f, t.g)

domain(t::ComposedTransformation) =
    composed_domain(RR_stability(t.f), RR_stability(t.g), t.f, t.g)

function logjac(t::ComposedTransformation, x)
    @unpack f, g = t
    y = g(x)
    log_g′x = logjac(g, x)
    log_f′y = logjac(f, y)
    log_f′y + log_g′x
end

inverse(t::ComposedTransformation, x) = inverse(t.g, inverse(t.f, x))

inverse(t::ComposedTransformation) =
    ComposedTransformation(inverse(t.g), inverse(t.f))

isincreasing(c::ComposedTransformation) = isincreasing(c.f) == isincreasing(c.g)

∘(f::UnivariateTransformation, g::UnivariateTransformation) =
    ComposedTransformation(f, g)

∘(f::Affine, g::Affine) = Affine(f.α*g.α, _fma(f.α, g.β, f.β))

# calculated transformations

"""
    affine_bridge(interval1, interval1)

Return an affine transformation between two intervals of the same type.
"""
function affine_bridge(x::Segment, y::Segment)
    α = width(y) / width(x)
    β = middle(y) - middle(x) * α
    Affine(α, β)
end

affine_bridge(::RealLine, ::RealLine) = IDENTITY

affine_bridge(x::PositiveRay, y::PositiveRay) = Affine(1, y.left - x.left)
affine_bridge(x::NegativeRay, y::NegativeRay) = Affine(1, y.right - x.right)
affine_bridge(x::PositiveRay, y::NegativeRay) =
    Affine(1, y.right + x.left) ∘ NEGATION
affine_bridge(x::NegativeRay, y::PositiveRay) =
    Affine(1, y.left + x.right) ∘ NEGATION

"""
    default_transformation(dom, img)

Return a transformation from `dom` that can be mapped to `img` using
`affine_bridge`.
"""
default_transformation(::RealLine, ::Segment) = LOGISTIC
default_transformation(::RealLine, ::PositiveRay) = EXP
default_transformation(::RealLine, ::NegativeRay) = EXP

default_transformation(::Segment, ::RealLine) = LOGIT
default_transformation(::PositiveRay, ::RealLine) = LOG
default_transformation(::NegativeRay, ::RealLine) = LOG

bridge(dom::RealLine, img, t = default_transformation(dom, img)) =
    affine_bridge(image(t), img) ∘ t

bridge(dom, img::RealLine, t = default_transformation(dom, img)) =
    t ∘ affine_bridge(dom, domain(t))

bridge(::RealLine, ::RealLine) = IDENTITY

# Special constants

"""
Hyperbolic tangent transformation.

An affine stretch of LOGISTIC to (-1, 1).
"""
const TANH = bridge(ℝ, Segment(-1.0, 1.0), LOGISTIC)

# UnitVector

"""
    UnitVector(n)

Transform `n-1` real numbers to a unit vector of length `n`, under the Euclidean
norm.
"""
struct UnitVector <: ContinuousTransformation
    n::Int
    function UnitVector(n::Int)
        @argcheck n ≥ 1 "Dimension should be positive."
        new(n)
    end
end

length(t::UnitVector) = t.n - 1

rhs_string(t::UnitVector, term) = "UnitVector($(t.n))($term)"

function transform_and_logjac!(ys::AbstractVector{T},
                               t::UnitVector, xs::AbstractVector) where T
    @unpack n = t
    @argcheck length(xs) == n - 1
    @argcheck length(ys) == n
    remainder = one(T)
    logjac = zero(T)
    for (i, x) in enumerate(xs)
        z = @increment_logjac logjac TANH(x)
        ys[i] = z * √remainder
        logjac += log(remainder) / 2
        remainder -= abs2(z) * remainder
    end
    ys[end] = √remainder
    ys, logjac
end

transform_and_logjac(t::UnitVector, zs::AbstractVector{T}) where T =
    transform_and_logjac!(Vector{T}(undef, t.n), t, zs)

function inverse!(xs::AbstractVector{T},
                  t::UnitVector, ys::AbstractVector) where T
    @unpack n = t
    @argcheck length(xs) == n - 1
    @argcheck length(ys) == n
    remainder = one(T)
    for i in 1:(n - 1)
        y = ys[i]
        z = y / √remainder
        xs[i] = inverse(TANH, z)
        remainder -= abs2(y)
    end
    xs
end

inverse(t::UnitVector, ys::AbstractVector{T}) where T =
    inverse!(Vector{T}(undef, t.n - 1), t, ys)

@define_from_transform_and_logjac UnitVector

# Cholesky factor of a correlation matrix

"""
    CorrelationCholeskyFactor(n)

Cholesky factor of a correlation matrix of size `n`.
"""
struct CorrelationCholeskyFactor <: ContinuousTransformation
    n::Int
    function CorrelationCholeskyFactor(n)
        @argcheck n ≥ 1 "Dimension should be positive."
        new(n)
    end
end

"""
    $SIGNATURES

Number of elements in a triangle of an ``n×n`` square matrix, including the
diagonal.
"""
triangle_length(n) = n * (n-1) ÷ 2

length(t::CorrelationCholeskyFactor) = triangle_length(t.n)

rhs_string(t::CorrelationCholeskyFactor, term) =
    "CorrelationCholeskyFactor($(t.n))($term)"

function transform_and_logjac(t::CorrelationCholeskyFactor,
                              x::AbstractVector{T}) where T
    @unpack n = t
    @argcheck length(x) == length(t)
    logjac = zero(T)
    cumulative_index = 0
    L = zeros(T, n, n)
    for i in 1:n
        k = i - 1
        _, lj = transform_and_logjac!(@view(L[i, 1:i]), UnitVector(i),
                                      @view(x[cumulative_index .+ (1:k)]))
        logjac += lj
        cumulative_index += k
    end
    LowerTriangular(L), logjac
end

@define_from_transform_and_logjac CorrelationCholeskyFactor

function inverse(t::CorrelationCholeskyFactor,
                 L::LowerTriangular{T}) where T
    @unpack n = t
    @argcheck size(L) == (n, n)
    x = Vector{T}(undef, triangle_length(n))
    cumulative_index = 0
    for i in 1:n
        k = i - 1
        inverse!(@view(x[cumulative_index .+ (1:k)]), UnitVector(i),
                 @view(L[i, 1:i]))
        cumulative_index += k
    end
    x
end

# helper function for Cholesky decompositions

"""
    $SIGNATURES

Log PDF of the LKJ distribution (Lewandowski et al 2009) for correlation
matrices.

A correlation matrix ``Ω=LL'`` has the density ``|Ω|^(η-1)``. However, it is
usually not necessary to construct ``Ω``, so this function is parametrized in
terms of a Cholesky decomposition `L`.

Note that this function **does not check if `L` yields a valid correlation
matrix**. Use [`CorrelationCholeskyFactor`](@ref) for generating valid
arguments.

Valid values are ``η > 0``. When ``η > 1``, the distribution is unimodal at the
identity, while ``0 < η < 1`` has a trough. ``η = 2`` is recommended as a vague
prior.

When ``η = 1``, the density is uniform. Note however that the function should
still be invoked, because of the Jacobian correction of the transformation.
"""
function lkj_correlation_cholesky_logpdf(L::Union{LowerTriangular,
                                                  UpperTriangular}, η)
    @argcheck η > 0
    z = diag(L)
    n = size(L, 1)
    sum(log.(z) .* ((n:-1:1) .+ 2*(η-1))) + log(2) * n
end

# grouped transformations

"""
$TYPEDEF

Abstract type for grouped transformations.

A grouped transformation takes a vector, and transforms contiguous blocks of
elements to some output type, determined by the specific transformation type.

All subtypes support

- `length`: return the length of the vector that can be used as an argument

- callable object for the transformation

- `logjac`, and `inverse`,

- `domain` and `image`, which may have specific interpretation for their result
  types depending on the concrete subtype.
"""
abstract type GroupedTransformation <: ContinuousTransformation end

"""
    $SIGNATURES

Return the transformation from a wrapper object.
"""
get_transformation(d) = d.transformation

# array transformations

"""
    ArrayTransformation(transformation, dims)
    ArrayTransformation(transformation, dims...)

Apply transformation to a vector, returning an array of the given dimensions.

[`domain`](@ref), [`image`](@ref), and [`isincreasing`](@ref) return the
corresponding values for the underlying transformation.
"""
struct ArrayTransformation{T <: UnivariateTransformation,
                           D} <: GroupedTransformation
    transformation::T
    function ArrayTransformation(transformation::T,
                                 dims::Tuple{Vararg{Int64, N}}) where {T,N}
        @argcheck all(dims .> 0) "Invalid dimensions."
        new{T,dims}(transformation)
    end
end

ArrayTransformation(transformation, dims::Int...) =
    ArrayTransformation(transformation, dims)

size(t::ArrayTransformation{T,D}) where {T,D} = D

length(t::ArrayTransformation) = prod(size(t))

function transformation_string(t::ArrayTransformation, term)
    s = size(t)
    s_print = length(s) == 1 ? s[1] : s
    "$(transformation_string(t.transformation, term)) for $(s_print) elements"
end

(t::ArrayTransformation)(x) = reshape((t.transformation).(x), size(t))

function logjac(t::ArrayTransformation, x)
    @argcheck length(x) == length(t) DimensionMismatch()
    lj(x) = logjac(t.transformation, x)
    sum(lj, x)
end

inverse(t::ArrayTransformation, x) =
    reshape(inverse.(t.transformation, x), size(t))

@forward ArrayTransformation.transformation isincreasing, image, domain

# transformation tuple

"""
    TransformationTuple(transformations::Tuple)
    TransformationTuple(transformations...)

A tuple of [`ContinuousTransformation`](@ref)s. Given a vector of matching
`length`, each takes as many reals as needed, and returns the result as a tuple.
"""
@auto_hash_equals struct
    TransformationTuple{T <: Tuple{Vararg{ContinuousTransformation}}} <:
        GroupedTransformation
    transformations::T
end

TransformationTuple(ts::ContinuousTransformation...) = TransformationTuple(ts)

function transformation_string(t::TransformationTuple, x)
    s = IOBuffer()
    print(s, "TransformationTuple")
    for (tt, ix) in zip(t.transformations,
                        transformation_indexes(t.transformations))
        print(s, "\n    ")
        print(s, transformation_string(tt, "$(x)[$(ix)]"))
    end
    String(take!(s))
end

length(t::TransformationTuple) = sum(length(t) for t in t.transformations)

domain(t::TransformationTuple) = domain.(t.transformations)

image(t::TransformationTuple) = image.(t.transformations)

getindex(t::TransformationTuple, ix) = t.transformations[ix]

function next_indexes(acc, t)
    l = length(t)
    acc + l, acc .+ (1:l)
end

next_indexes(acc, t::UnivariateTransformation) = acc+1, acc+1

#=@unroll=# function transformation_indexes(ts)
    acc = 0
    result = ()
    #=@unroll=# for t in ts
        acc, ix = next_indexes(acc, t)
        result = (result..., ix)
    end
    result
end

function (t::TransformationTuple)(x)
    ts = t.transformations
    map((t,ix) -> t(x[ix]), ts, transformation_indexes(ts))
end

function logjac(t::TransformationTuple, x)
    ts = t.transformations
    sum(map((t,ix) -> logjac(t, x[ix]), ts, transformation_indexes(ts)))
end

inverse(t::TransformationTuple, y::Tuple) =
    vcat(map(inverse, t.transformations, y)...)

# wrapper

"""
$TYPEDEF

Wrap a transformation to achieve some specialized functionality.

Supports `length`, [`get_transformation`](@ref), and other methods depending on
the subtype.
"""
abstract type TransformationWrapper <: Function end

# loglikelihood wrapper

"""
    TransformLogLikelihood(ℓ, transformation::Union{Tuple, GroupedTransformation})

    TransformLogLikelihood(ℓ, transformations...)

Return a callable that

1. transforms its vector argument using a grouped transformation to a set of
   values,

2. calls `ℓ` (which should return a scalar) with this tuple.

3. returns the result above corrected by the log Jacobians.

Useful when `ℓ` is a log-likelihood function with a restricted domain, and
`transformations` is used to trasform to this domain from ``ℝ^n``.

See also [`get_transformation`](@ref), [`get_distribution`](@ref),
`Distributions.logpdf`, and [`logpdf_in_domain`](@ref).
"""
struct TransformLogLikelihood{L, T <: GroupedTransformation} <:
        TransformationWrapper
    loglikelihood::L
    transformation::T
end

TransformLogLikelihood(L, T::Tuple) =
    TransformLogLikelihood(L, TransformationTuple(T))

TransformLogLikelihood(L, ts::ContinuousTransformation...) =
    TransformLogLikelihood(L, TransformationTuple(ts))

"""
    $SIGNATURES

Return the log likelihood function.
"""
get_loglikelihood(t::TransformLogLikelihood) = t.loglikelihood

@forward TransformLogLikelihood.transformation length

(f::TransformLogLikelihood)(x) =
    f.loglikelihood(f.transformation(x)) + logjac(f.transformation, x)

function show(io::IO, f::TransformLogLikelihood)
    print(io, "TransformLogLikelihood of length $(length(f)), with ")
    print(io, f.transformation)
end

# transform distributions

"""
    TransformDistribution(distribution, transformation)

Given a `transformation` and a `distribution`, create a transformed distribution
object that has the distribution of `transformation(x)` with `x ∼ distribution`.

The transformation object is callable with the same syntax as
`transformation`. It also supports methods `rand`, `length`.

See also [`logpdf_in_domain`](@ref) and [`logpdf_in_image`](@ref).
"""
struct TransformDistribution{D <: Union{ContinuousUnivariateDistribution,
                                        ContinuousMultivariateDistribution},
                             T <: Union{UnivariateTransformation,
                                        GroupedTransformation}} <:
                                            TransformationWrapper
    distribution::D
    transformation::T
    function TransformDistribution(distribution::D, transformation::T) where
        {D <: ContinuousUnivariateDistribution, T <: UnivariateTransformation}
        new{D, T}(distribution, transformation)
    end
    function TransformDistribution(distribution::D, transformation::T) where
        {D <: ContinuousMultivariateDistribution, T <: GroupedTransformation}
        @argcheck(length(distribution) == length(transformation),
                  "Incompatible lengths.")
        new{D,T}(distribution, transformation)
    end
end

get_transformation(t::TransformDistribution) = t.transformation

"""
    $SIGNATURES

Return the wrapped distribution.
"""
get_distribution(t::TransformDistribution) = t.distribution

(t::TransformDistribution)(x) = t.transformation(x)

rand(t::TransformDistribution) = t(rand(t.distribution))

@forward TransformDistribution.distribution length

"""
    $SIGNATURES

The log pdf for a transformed distribution at `t(x)` in image, calculated in the
domain without performing the transformation.

The log pdf is adjusted with the log determinant of the Jacobian, ie the
following holds:

```julia logpdf_in_image(t, t(x)) == logpdf_in_domain(t, x) ```

See [`logpdf_in_image`](@ref).

!!! note

    Typical usage of this function would be drawing some random `x`s from the
    contained distribution (possibly also used for some other purpose), and
    obtaining the log pdfs at `t(y)` with the same values.
"""
function logpdf_in_domain(t::TransformDistribution, x)
    # NOTE adjusting by -logjac because of the derivative of the inverse rule
    logpdf(t.distribution, x) - logjac(t.transformation, x)
end

"""
    $SIGNATURES

The log pdf for a transformed distribution at `y` in image.

See also [`logpdf_in_domain`](@ref).
"""
function logpdf_in_image(t::TransformDistribution, y)
    x = inverse(t.transformation, y)
    logpdf_in_domain(t, x)
end

# misc utilities
# NOTE: may be experimental, removed/redesigned
# NOTE: code written for simplicity, not speed, benchmark & optimize later

"""
    $SIGNATURES

Make a container for ungrouped results.

The first argument, `Vector` or `Array`, determines the result of ungrouping.

`representative_elt` is a representative element, used to determine element
type and stucture, and `len` is the length of the result.
"""
function make_ungrouped_container(::Type{Vector},
                                  representative_elt::AbstractArray{T},
                                  len) where T
    [Vector{T}(undef, len) for _ in CartesianIndices(axes(representative_elt))]
end

function make_ungrouped_container(::Type{Array},
                                  representative_elt::AbstractArray{T},
                                  len) where T
    Array{T}(undef, len, size(representative_elt)...)
end

make_ungrouped_container(T, representative_elt::Tuple, len) =
    map(x -> make_ungrouped_container(T, x, len), representative_elt)

make_ungrouped_container(_, representative_elt::T, len) where {T <: Real} =
    Vector{T}(undef, len)

"""
    $SIGNATURES

Save `elt` in `container` at index `i`. `T` determines the layout of
`container`, conformably with [`make_ungrouped_container`](@ref).
"""
ungroup_elt!(T, container::NTuple{N, Any}, elt::NTuple{N, Any}, i) where N =
    foreach((c, e) -> ungroup_elt!(T, c, e, i), container, elt)

function ungroup_elt!(::Type{Vector}, container::AbstractArray,
                      elt::AbstractArray, i)
    @argcheck axes(container) == axes(elt)
    for j in 1:length(elt)
        ungroup_elt!(Vector, container[j], elt[j], i)
    end
end

ungroup_elt!(::Type{Array}, container::AbstractArray,
             elt::AbstractArray{T, N}, i) where {T, N} =
    container[i, fill(Colon(), N)...] .= elt

ungroup_elt!(_, container::AbstractVector, elt::Number, i) = container[i] = elt

"""
    $SIGNATURES

Map the vector `A` elementwise using `f`, and collect the results in arrays or
vectors organized as a single result from `f`. For example, if `f` returns
tuples, this function returns tuples of vectors.

The shape of the result is determined by the first argument, which can be
`Vector` or `Array`. The difference is that `Array` stacks arrays into arrays,
while `Vector` does not.

`Array` may be more useful for summaries along dimensions, eg `mean` and
similar.
"""
function ungrouping_map(T, f, A::AbstractVector)
    n = length(A)
    representative_elt = f(A[1])
    container = make_ungrouped_container(T, representative_elt, n)
    ungroup_elt!(T, container, representative_elt, 1)
    for i in 2:n
        ungroup_elt!(T, container, f(A[i]), i)
    end
    container
end

"""
    $SIGNATURES

Test if a numerical argument is made up of finite real numbers.
"""
all_finite(x::Real) = isfinite(x)

all_finite(A::AbstractArray{T}) where {T <: Real} = all(isfinite, A)

all_finite(t::Tuple) = all(all_finite, t)

all_finite(itr) = all(all_finite, itr)

end


using ModelingToolkit
using ModelingToolkit: value
using Test

@parameters t
@variables x(t) y(t) z(t)

null_op = 0*t
@test isequal(simplify(null_op), 0)

one_op = 1*t
@test isequal(simplify(one_op), t)

identity_op = Num(Term(identity,[x.val]))
@test isequal(simplify(identity_op), x)

minus_op = -x
@test isequal(simplify(minus_op), -1x)
simplify(minus_op)

@variables x

@test toexpr(expand_derivatives(Differential(x)((x-2)^2))) == :($(*)(2, $(+)(-2, x)))
@test toexpr(expand_derivatives(Differential(x)((x-2)^3))) == :($(*)(3, $(^)($(+)(-2, x), 2)))
@test toexpr(simplify(x+2+3)) == :($(+)(5, x))

d1 = Differential(x)((-2 + x)^2)
d2 = Differential(x)(d1)
d3 = Differential(x)(d2)

@test toexpr(expand_derivatives(d3)) == :(0)
@test toexpr(simplify(x^0)) == :(1)

@test ModelingToolkit.substitute(value(2x + y == 1), Dict(x => 0.0, y => 0.0)) === false
@test ModelingToolkit.substitute(value(2x + y == 1), Dict(x => 0.0, y => 1.0)) === true

# 699
using SymbolicUtils: substitute
@parameters t a(t) b(t)

# back and forth substitution does not work for parameters with dependencies
term = value(a)
term2 = substitute(term, a=>b)
@test term2 isa Term{ModelingToolkit.Parameter{Real}}
@test isequal(term2, b)
term3 = substitute(term2, b=>a)
@test term3 isa Term{ModelingToolkit.Parameter{Real}}
@test isequal(term3, a)


using EchogramUtils
using Test

@test pow2db(1000) == 30
@test db2pow(30) == 1000

@test mag2db(10) == 20
@test db2mag(20) == 10


using ..Statistics
using ..Bases

typealias GPCDistribution ContinuousUnivariateDistribution

"A gpc pair, capturing the pairing of distributions with systems of orthogonal polymials"
immutable GPCPair
  dist::GPCDistribution
  polysys::PolynomialSystem
end

"The standard mapping from distribution to orthogonal polynomials"
map_dist_to_polysys = Dict{GPCDistribution, PolynomialSystem}()
map_dist_to_polysys[Uniform(-1,1)] = LegendrePolynomials()
map_dist_to_polysys[Normal(0,1)] = HermitePolynomials()

"Constructor for GPCPairs "
GPCPair(dist) = GPCPair(dist, map_dist_to_polysys[dist])

"The abstract base for gpc germs (maybe we extend that later to infinite germs)"
abstract GPCGermBase

"A gpc germ of fixed length (the standard case)"
immutable GPCGerm <: GPCGermBase
  pairs::Array{GPCPair}
end
function GPCGerm(dist::GPCDistribution, n::Int)
  GPCGerm(fill(GPCPair(dist), n))
end

import Base.length
length(germ::GPCGerm) = length(germ.pairs)

distributions(germ::GPCGerm) = [pair.dist for pair in germ.pairs]

import Base.rand
rand(germ::GPCGerm) = Float64[rand(dist) for dist=distributions(germ)]

import Distributions.pdf
function pdf{T<:Real}(germ::GPCGerm, x::Vector{T})
  prod([pdf(dist, xi) for (dist,xi) in zip(distributions(germ), x)])
end

import Distributions.cdf
function cdf{T<:Real}(germ::GPCGerm, x::Vector{T})
  prod([cdf(dist, xi) for (dist,xi) in zip(distributions(germ), x)])
end

export GPCGerm,
       GPCPair


"""
function of the type {0, .., k-1} -> {0, ..., l-1}
represened internally as array of integers
"""
struct ResidueFunction <: BasicDiscreteFunction
    m :: Vector{Int}
    dom :: Segment{Int}
    rng :: Segment{Int}
end

"""
function of the type
ℤ_k1 × … × ℤ_km → ℤ_l
represened internally as array of integers
"""
struct ExtendedResidueFunction{T, N} <: BasicDiscreteFunction
    m :: Vector{Int}
    dom :: DirectProduct{T}
    rng :: Segment{Int}
    sizes :: NTuple{N, Int}
end

const ResFun = ResidueFunction
const ExtFun = ExtendedResidueFunction

# ResidueFunction from domain and codomain
function ResFun(dom :: Segment, codom :: Segment)
    m = zeros(Int, length(dom))
    return ResFun(m, dom, codom)
end
# constructing [0..k-1] -> [0..l-1]
# input : max_x, max_y
ResFun(k, l) = ResFun(Segment(k), Segment(l))

# from tuple of args
# input : (tuple) = (max_x, max_y)
ResFun(tup :: Tuple{T, S}) where {T, S <: Integer} = ResFun(tup[1], tup[2])

# constructiong {0, ..., k-1} -> {0, .., k-1}
ResFun(k) = ResFun(k, k)

# ExtFun from domain and codomain
function ExtFun(dom :: AbstractProduct, codom :: Segment)
    m = zeros(Int, length(dom))
    return ExtFun(m, dom, codom, size(dom))
end

# ExtFun from tuple (sizes) and max_y
function ExtFun(sizes :: NTuple, max_y :: Int)
    dom = DirectProduct(map(Segment, sizes)...)
    codom = Segment(max_y)
    return ExtFun(dom, codom)
end

(f :: ResFun)(x :: Integer) = f.m[x + 1]
Base.setindex!(f :: ResFun, value, key) = setindex!(f.m, value, key + 1)

function (f :: ExtFun)(x :: Union{NTuple, Vector})
    ind = itr2ind(x, f.sizes)
    return f.m[ind + 1]
end

function Base.setindex!(f :: ExtFun, v, k :: NTuple)
    ind = itr2ind(k, f.sizes)
    f.m[ind+1] = v
end
Base.setindex!(f :: ExtFun, v, k...) = setindex!(f, v, tuple(k...))


## Intro figures
using Plots
using GraphRecipes
using NearestNeighbors
using Random
using SparseArrays
using LinearAlgebra
## Load the graph data
ei,ej = copy.(eachcol(Int.(readdlm("undir-graph-14.edges"))))
xy = readdlm("undir-graph-14.xy")'
n = size(xy,2)
A = sparse(ei,ej,1,n,n)
A = A+A'
##
pyplot(size=(300,300))
p = graphplot(ei, ej, x =xy[1,:], y=xy[2,:],
  markercolor=:white, markerstrokecolor=:black,
  markersize=8, linecolor=3, linealpha=1.0, linewidth=1.0,
  axis_buffer=0.02, background=:transparent,dpi=300)
savefig("undir-graph-14.pdf")
##
Matrix(A)
##
using MatrixNetworks
f,lam = fiedler_vector(A)
p = graphplot(ei, ej, x =xy[1,:], y=xy[2,:],
  markercolor=:white, markerstrokecolor=:black,
  markersize=1, linecolor=3, linealpha=1.0, linewidth=1.0,
  axis_buffer=0.02, background=:transparent,dpi=300)
scatter!(xy[1,:], xy[2,:], marker_z=f, markersize=8)

## Seeded PageRank vectors
# I tried Fiedler vectors on this graph, but they just pick off little
# pieces like usual :(
names = readlines("../2-unit-1-demos/artistsim.names")
data = readdlm("../2-unit-1-demos/artistsim.smat")
xy = readdlm("artistsim.xy")
##
A = sparse(Int.(data[2:end,1]).+1,
           Int.(data[2:end,2]).+1,
           Int.(data[2:end,3]),
           Int(data[1,1]),Int(data[1,1]))
Asym = max.(A,A') # make the graph undirected
##
Acc = largest_component(Asym)[1]
##
#f,lam = fiedler_vector(Acc)
##
scatter(xy[:,1],xy[:,2], marker_z=(f),colorbar=true,alpha=0.5,
  markerstrokecolor=nothing,markerstrokewidth=0)
##
ei,ej = findnz(triu(Asym,1))
graphplot(ei, ej, x =xy[:,1], y=xy[:,2],
  markercolor=:black, markerstrokecolor=:white,
  size=(1200,1200),dpi=300,
  markersize=2, linecolor=1, linealpha=0.01, linewidth=0.5,
  markeralpha=0.2,
  marker_z = f,
  axis_buffer=0.02, background=nothing)
##
f = seeded_pagerank(Asym, 0.85, 17998)
##
graphplot(ei, ej, x =xy[:,1], y=xy[:,2],
  markercolor=:black, markerstrokecolor=:white,
  size=(1200,1200),dpi=300,
  markersize=2, linecolor=1, linealpha=0.01, linewidth=0.5,
  markeralpha=0.2,
  marker_z = log10.(f.+eps(1.0)),
  axis_buffer=0.02, background=nothing)
##
#f2 = seeded_pagerank(Acc, 0.5, 17998)
f2 = seeded_pagerank(Acc, 0.85, 17998)
p = sortperm(f2) # this gets things to draw in the right z-order
scatter(xy[p,1],xy[p,2], marker_z=log10.(f2[p]),colorbar=false,alpha=0.5,
  markerstrokecolor=nothing,markerstrokewidth=0,label="",framestyle=:none,background=nothing)
##
f2 = seeded_pagerank(Acc, 0.85, 200)
p = sortperm(f2) # this gets things to draw in the right z-order
scatter(xy[p,1],xy[p,2], marker_z=-abs.(log10.(f2[p])).^0.5,colorbar=false,alpha=0.5,
  color=:magma,
  markerstrokecolor=nothing,markerstrokewidth=0,label="",framestyle=:none,background=nothing)
##
savefig("artistsim-spr-200.pdf")


# # "Cosine advection-diffusion" Spatial resolution convergence test

using PyPlot

using Oceananigans.Grids

# Define a few utilities for running tests and unpacking and plotting results

include("ConvergenceTests/ConvergenceTests.jl")

run_test = ConvergenceTests.OneDimensionalCosineAdvectionDiffusion.run_test

""" Run advection-diffusion test for all Nx in resolutions. """
function run_convergence_test(κ, U, resolutions...)

    # Determine save time-step
             Lx = 2π
      stop_time = 0.01
              h = Lx / maximum(resolutions)
    proposal_Δt = 1e-3 * min(h / U, h^2 / κ)

    # Adjust time-step
    stop_iteration = round(Int, stop_time / proposal_Δt)
                Δt = stop_time / stop_iteration

    # Run the tests
    results = [run_test(Nx=Nx, Δt=Δt, stop_iteration=stop_iteration, U=U, κ=κ) for Nx in resolutions]

    return results
end

#####
##### Run test
#####

Nx = 2 .^ (3:7) # N = 64 through N = 256
diffusion_results = run_convergence_test(1e-1, 0, Nx...)
advection_results = run_convergence_test(1e-6, 3, Nx...)
advection_diffusion_results = run_convergence_test(1e-2, 1, Nx...)

#####
##### Plot solution and error profile
#####

all_results = (diffusion_results, advection_results, advection_diffusion_results)
names = ("diffusion only",  "advection only",  "advection-diffusion")
linestyles = ("-", "--", ":")
specialcolors = ("xkcd:black", "xkcd:indigo", "xkcd:wine red")

# Solutions
close("all")

fig, axs = subplots(nrows=2, figsize=(12, 6), sharex=true)

ConvergenceTests.OneDimensionalUtils.plot_solutions!(axs, all_results, names, linestyles, specialcolors)

savefig("figs/cosine_advection_diffusion_solutions.png", dpi=480)

# Error profile
fig, axs = subplots()

ConvergenceTests.OneDimensionalUtils.plot_error_convergence!(axs, Nx, all_results, names)

savefig("figs/cosine_advection_diffusion_error_convergence.png", dpi=480)


# Using Convex.jl to test if the probability distributions in Homework problem 2
# Sets up the optimization problem for the Cut Inflation for the triangle causal
# model with the distribution:
#
# P(a, b, c; t) = if (a == b == c) => t/2
#                 else             => (1 - t)/6
#
# the optimal value of t is found through repetition

using Convex
using GLPK
using Inflation
using Printf

struct InflationLoopData
	all_feasible::Bool
	all_infeasible::Bool
	tmin::Float64
	tmax::Float64
	Ndiv::Int64
end

function check_inflation_t(infl_strings::Vector{String}, M::Array{Int64, 2}, idata::InflationLoopData)::InflationLoopData
	# checks if an inflation is feasible for pfunct, for certain values of t
	tmin = idata.tmin
	tmax = idata.tmax
	Ndiv = idata.Ndiv
	dt   = (tmax - tmin)/(Ndiv - 1)
	for (i, t) in enumerate(range(tmin, tmax, length = Ndiv))
		# avoid repeatedly testing the 1st element
		if i == 1 && t != 0.0
			continue
		end

		# get pfunc(a, b, c; t) and get the associated b-array
		pfunct = get_pfunct(t)
		b      = get_b_array(infl_strings, pfunct, cut)

		# the variables that get optimized by Convex.jl
		v = get_v_array_Variable(infl_strings)

		# set up the objective and constraints
		objective   = 0
		constraints = [v >= 0, M*v == b]

		# create the problem and solve
		problem = Convex.maximize(objective, constraints)
		solve!(problem, GLPK.Optimizer, verbose=false)

		# return the results
		if problem.status == Convex.MathOptInterface.INFEASIBLE
			@printf("Checked t = %0.6f : solution is infeasible.\n", t)
			if (i == 1) # all infeasible
				return InflationLoopData(false, true, tmin, tmax, Ndiv)
			else
				new_tmin = t - (tmax - tmin)/(Ndiv - 1)
				new_tmax = t
				return InflationLoopData(false, false, new_tmin, new_tmax, Ndiv)
			end
		else
			@printf("Checked t = %0.6f : solution is feasible.\n", t)
		end
	end

	return InflationLoopData(true, false, tmin, tmax, Ndiv)
end

# the probability matrix is set up as
# P_{A2,B1} = sum_{C1}P_{A2,B1,C1}   --> 'yyx'  (3rd term is integrated)
# P_{A2,C1} = sum_{B1}P_{A2,B1,C1}   --> 'yxy'  (2nd term is integrated)
# P_{B1,C1} = sum_{A2}P_{A2,B1,C1}   --> 'xyy'  (1st term is integrated)
infl_strings = ["yyx", "yxy", "xyy"]

# creating the M-matrix to be used in the Convex.jl optimization
M = get_M_matrix(infl_strings)

tmin = 0.0  # smallest allowed value of t
tmax = 1.0  # greatest allowed value of t
Ndiv = 11   # number of elements to divide the range into
idata = InflationLoopData(false, false, tmin, tmax, Ndiv)

# check to 5 decimal places
for i in 1:6
	global idata
	idata = check_inflation_t(infl_strings, M, idata)
	if idata.all_feasible
		@printf("All values of t between %0.6f and %0.6f are feasible.\n", tmin, idata.tmax)
		break
	elseif idata.all_infeasible
		@printf("Unable to find further feasible values.\n")
		break
	else
		idata = InflationLoopData(false, false, idata.tmin, idata.tmax, idata.Ndiv)
	end
end

@printf("All values of t between %0.6f and %0.6f are feasible.\n", tmin, idata.tmin)