abstract type  AbstractMaxEntropyModel   end

# Large version UBCM
struct UBCM{T} <: AbstractMaxEntropyModel where {T<:Real}
    x0::Vector{T}         # initial solution estimate
    x::Vector{T}          # solution if converged
    F::Vector{T}          # vector for values of the system of equations
    f!::Function          # sytem that needs to be solved
end

show(io::IO, model::UBCM{T}) where {T} = print(io, "UBCM ($(T)) model ($(length(model.x0)) nodes)")

function UBCM(G::T) where {T<:AbstractGraph}
    # assert validity of the method
    !is_connected(G) && throw(DomainError("Graph is not connected"))
    has_self_loops(G) && throw(DomainError("Graph has self loops"))
    # get degree sequence
    k = degree(G)
    # estimate for x0
    x0 = k ./ maximum(k)
    return UBCM(x0, similar(x0), similar(x0), (F::Vector, x::Vector) -> UBCM!(F, x, k))
end

function UBCM!(F::Vector, x::Vector, k::Vector)
    Threads.@threads for i in eachindex(x)
        @inbounds F[i] = -f(x[i]*x[i]) - k[i]
        @simd for j in eachindex(x)
            @inbounds F[i] += f(x[i]*x[j])
        end
    end
    return F
end

function solve!(model::UBCM)
    res =  nlsolve(model.f!, model.x0, autodiff=:forward)
    !res.f_converged && @warn "Solver did not converge"
    if res.f_converged
        any(res.zero .< 0) && @warn "Negative values in solution"
        @inbounds for i in eachindex(model.x)
            model.x[i] = res.zero[i]
        end
    end
    return res
end

# Compact version UBCM
mutable struct UBCMCompact{T, I} <: AbstractMaxEntropyModel where {T<:Real, I<:Integer}
    x0::Vector{T}         # initial solution estimate
    xs::Vector{T}         # solution if converged
    x::Union{IndirectArray, Nothing}      # expanded vector (using IndirectArrays)
    F::Vector{T}          # vector for values of the system of equations
    f!::Function          # sytem that needs to be solved
    #freqs::Dict{Int, Int} # dict holding the mapping (not required, but nice to have)
    #nodemap::Dict{Int, Int} # dict holding the mapping
    revmap::Vector{I}        # dict holding the reverse mapping
end

show(io::IO, model::UBCMCompact{T,I}) where {T,I} = print(io, "Compact UBCM ($(T),$(I)) model ($(length(model.revmap)) nodes reduced to $(length(model.x0)))")

function UBCMCompact(G::T) where {T<:AbstractGraph}
    # assert validity of the method
    !is_connected(G) && throw(DomainError("Graph is not connected"))
    has_self_loops(G) && throw(DomainError("Graph has self loops"))
    # get unique degrees and frequencies
    d = degree(G)
    freqs = countmap(d)
    κ = collect(keys(freqs))
    # get reverse mapping
    #revmap = Dict(i => findall(x->isequal(x, κ[i]), d) for i in eachindex(κ))
    #nodemap = Dict(keys => val for (key, val) in revmap)
    indtype = length(d) < typemax(UInt8) ? UInt8 : length(d) < typemax(UInt16) ? UInt16 : length(d) < typemax(UInt32) ? UInt32 : UInt64
    revmap = Vector{indtype}(undef, length(d))
	for i in eachindex(κ)
		revmap[findall(x->isequal(x, κ[i]), d)] .= indtype(i) # link values of large vector to compact one
	end

    # estimate for x0
    x0 = sqrt.(κ) ./ sum(κ)

    return UBCMCompact(x0, similar(x0), nothing, similar(x0), (F::Vector, x::Vector) -> UBCMCompact!(F, x, κ, freqs), revmap)
end

function UBCMCompact!(F::Vector, x::Vector, κ::Vector, freqs::Dict)
    Threads.@threads for i in eachindex(x)
        @inbounds F[i] = -f(x[i]*x[i]) - κ[i]
        @simd for j in eachindex(x)
            @inbounds F[i] += freqs[κ[j]]*f(x[i]*x[j])
        end
    end
    return F
end

function solve!(model::UBCMCompact)
    res =  nlsolve(model.f!, model.x0, autodiff=:forward)
    !res.f_converged && @warn "Solver did not converge"
    if res.f_converged
        # expansion of compact solution to large one happens here
        any(res.zero .< 0) && @warn "Negative values in solution"
        @inbounds for i in eachindex(model.x0)
            model.xs[i] = res.zero[i]
        end
        # generation of large solution happens here (using IndirectArrays)
        model.x = IndirectArray(model.revmap, model.xs)
    end
    return res
end

# supporting functions for computations of UBCM
f(x) = x / (1 + x) # helper function
a(i::Int,j::Int,x::T) where {T<:AbstractVector} = @inbounds f(x[i]*x[j]) # single element of adjacency matrix
"""Obtain full adjacency matrix from a model"""
function adjacency_matrix(model::UBCM)
    A = zeros(length(model.x), length(model.x))
    for i in eachindex(model.x)
        for j in eachindex(model.x)
            if i≠j
                A[i,j] = f(model.x[i]*model.x[j])
            end
        end
    end
    return A
end

"""helper function for adjacency matrix"""
function indmap(model::UBCMCompact,T::DataType,i,j)
    return model.revmap[i] <= model.revmap[j] ? T(model.revmap[i]+model.revmap[j]*(model.revmap[j]-1)÷2) : T(model.revmap[j]+model.revmap[i]*(model.revmap[i]-1)÷2)
end

"""Obtain full adjacency matrix from a compact model (as an IndirectArray)"""
function adjacency_matrix(model::UBCMCompact)
    # generate vector of all possible unique values in the adjacency matrix
    values = @inbounds [f(model.xs[i]*model.xs[j]) for j = 1:length(model.xs) for i = 1:j]
    push!(values,zero(eltype(values))) # include the zero for later
    # infer storage index type
    T = length(values) < typemax(UInt8) ? UInt8 : length(values) < typemax(UInt16) ? UInt16 : UInt32
    # generate indices in matrix from xs
    indices = [indmap(model,T,i,j) for i in 1:length(model.revmap), j in 1:length(model.revmap)]
    # set diagonal to zero
    indices[diagind(indices)] .= T(length(values))
    # return indirectarray
    return IndirectArray(indices, values)
end


using LsqFit
using StatsBase

#####
# Multiple ways to estimate an autocorrelation time (tau).
#####
"""
Integrated autocorrelation time. Defined as `tau_int = sum(autocor(M))+0.5`.
"""
function tau_integrated(X::Vector{T}) where T<:Real
    return sum(autocor(X))+0.5
end

"""
Autocorrelation time obtained from fitting autocorrelation C(t) to `exp(-t/tau_fitting)`.
"""
function tau_fitting(X::Vector{T}) where T<:Real
    logautocor = T[]
    Cx = autocor(X)

    firstneg = findfirst(x->x<=0, Cx)
    logCx = firstneg != 0 ? log.(Cx[1:firstneg-1]) : log.(Cx)
    length(logCx) <= 2 ? warn("Fitting log autocorrelation with less than 3 datapoints.") : nothing

    # Fitting
    model(x,p) = p[1].*x+p[2]
    pstart = [-1., 1.]
    fitres = curve_fit(model, 1.0:1.0:length(logCx), logCx, pstart)

    return -1/fitres.param[1]
end

"""
Autocorrelation time obtained from either static or dynamic binning analysis.

Optional keyword `binsize`. Decides wether we do static (fixed, finite binsize)
or dynamic (varying binsize) binning. (Based on Ch. 3.4)
"""
function tau_binning(X::Vector{T}; binsize=0) where T<:Real

    if binsize == 0
        #dynamic binning with varying binsize
        bss, R = R_function(X)
        length(R) > 0 || error("Binning failed. Maybe min_nbins to large?")
        Rplateau = R[end]
        for i in 1:length(bss)-2
            width = abs(maximum(R[i:i+2]) - minimum(R[i:i+2]))
            if width <= 0.05
                Rplateau = R[i]
                break;
            end
        end
        Rplateau != R[end] || warn("Couldn't find plateau. Maybe not converged?")
    else
        #static binning with fixed, single binsize
        Rplateau = R_value(X, binsize)
    end

    return 0.5*(Rplateau-1)
end

#####
# Statistical error (effective standard deviation) calculated from explicit estimate of autocorrelation time
#####
## CORRECT?

"""
Statistical error for correlated data from tau_binning analysis (find plateau (default) or fixed bin size). (Eq. 3.14)
"""
tau_binning_error(X::Vector{T}; binsize=0) where T<:Real = sqrt(var(X)*(1+2*tau_binning(X, binsize=binsize))) # 1/length(X) factor "left out" to have tau_binning_error = std for uncorrelated data
tau_binning_error(X::Vector{T}; binsize=0) where T<:Complex = sqrt(tau_binning_error(real(X), binsize=binsize)^2 + tau_binning_error(imag(X), binsize=binsize)^2)
tau_binning_error(X::Array{T}; binsize=0) where T<:Number = squeeze(mapslices(ts->tau_binning_error(ts, binsize=binsize), X, ndims(X)), ndims(X))

Neff_binning(X::Vector{T}; binsize=0) where T<:Real = floor(Int, length(X)/(1+2*tau_binning(X, binsize=binsize)))

# function tau_binning_error(X::Array{T}; binsize=0) where T<:Real
#     N = size(X, ndims(X)) # length of time series
#     linX = reshape(X, (:, N))
#     Nel = size(linX, 1) # number of elements/time series
#     el_shape = size(X)[1:end-1]

#     errs = zeros(T, Nel)
#     for i in 1:size(errs, 1)
#         errs[i] = tau_binning_error(linX[i,:], binsize=binsize)
#     end
#     return reshape(errs, el_shape)
# end

"""
Statistical error for correlated data from integrated autocorrelation time. (Eq. 3.14)
"""
tau_integrated_error(X::Vector{T}) where T<:Real = sqrt(var(X) *(2*tau_integrated(X)))
tau_integrated_error(X::Vector{T}) where T<:Complex = sqrt(tau_integrated_error(real(X))^2 + tau_integrated_error(imag(X))^2)
tau_integrated_error(X::Array{T}) where T<:Number = squeeze(mapslices(ts->tau_integrated_error(ts), X, ndims(X)), ndims(X))

Neff_integrated(X::Vector{T}) where T<:Real = floor(Int, length(X)/(2*tau_integrated(X)))

"""
Statistical error for correlated data from fitting `C(t) ~ exp(-t/tau)`. (Eq. 3.14)
"""
tau_fitting_error(X::Vector{T}) where T<:Real = sqrt(var(X)*(1+2*tau_fitting(X)))
tau_fitting_error(X::Vector{T}) where T<:Complex = sqrt(tau_fitting_error(real(X))^2 + tau_fitting_error(imag(X))^2)
tau_fitting_error(X::Array{T}) where T<:Number = squeeze(mapslices(ts->tau_fitting_error(ts), X, ndims(X)), ndims(X))

Neff_fitting(X::Vector{T}) where T<:Real = floor(Int, length(X)/(1+2*tau_fitting(X)))


include("structs.jl")
include("physics.jl")
include("computational.jl")

Broadcast.broadcastable(eos::TaitEoS) = Ref(eos)
Broadcast.BroadcastStyle(::Type{<:TaitEoS}) =
                        Broadcast.DefaultArrayStyle{0}()

Broadcast.broadcastable(eos::IdealGasEoS) = Ref(eos)
Broadcast.BroadcastStyle(::Type{<:IdealGasEoS}) =
                        Broadcast.DefaultArrayStyle{0}()

# -------------------- Density --------------------------- #

## The tait and ideal gas equations helps us determine
## ρ̂ᵢ= mᵢ/ sᵢV , but we need ρᵢ = mᵢ / V because ρᵢ is
## used in continuity and Darcy's equations. The first
## two functions return ρ̂ᵢ, while the third one returns ρᵢ

function density(x::TaitEoS, P::AbstractFloat)
    if x.B + P < 0
        println("Less than 0: ", P)
    end
    return x.ρ₀ / (1 - x.C * log10((x.B + P)
                                   / (x.B + x.P₀)))
end

function density(x::IdealGasEoS, P::AbstractFloat)
    return P * x.M / R / x.T
end

function density!(sys::System)
    @. @views sys.ρ[:, :, 1] = density(ideal_gas, sys.P) *
                                        sys.s[:, :, 1]
    @. @views sys.ρ[:, :, 2] = density(tait_C₅H₁₂, sys.P) *
                                        sys.s[:, :, 2]
end
# -------------------------------------------------------- #

# -------------------- Pressure -------------------------- #

function find_pressure(; root_finder::Function, ρ₁::T, ρ₂::T,
        tait::TaitEoS, igas::IdealGasEoS,
        p::Parameters, niter::Integer=50,
        eps_x::T=1e-6)::Tuple{T, T} where T<:AbstractFloat

    function f(P::T)::T where T<:AbstractFloat
        ρ̂₁ = density(igas, P)
        s  = ρ₁ / ρ̂₁
        ρ̂₂ = ρ₂ / (1 - s)
        return (ρ̂₂ - tait.ρ₀) / ρ̂₂ - tait.C *
            log10((tait.B + P) / (tait.B + tait.P₀))
    end

    function fder(P::T)::T where T<:AbstractFloat
        ρ̂₁ = density(igas, P)
        s  = ρ₁ / ρ̂₁
        ρ̂₂ = ρ₂ / (1 - s)
        return -tait.ρ₀ / ρ̂₂^2  * R * igas.T * igas.M *
            ρ₁ * ρ₂ / (igas.M * P - R * igas.T * ρ₁)^2 -
            tait.C / (tait.B + P) / log(10)
    end

    let
        P  = root_finder(; f=f, fder=fder, eps_x=eps_x)
        ρ̂₁ = density(igas, P)
        s  = ρ₁ / ρ̂₁

        return P, s
    end
end

function find_pressure!(sys::System, p::Parameters)
    for index in CartesianIndices(sys.P)
        i, j = Tuple(index)
        ## This gives seg fault, but is type stable
        # (sys.P[i, j], sys.s[i, j])::Tuple{T, T} =
        #             find_pressure(p=p,
        #             root_finder=newton_raphson,
        #             ρ₁=sys.ρ[i, j, 1],
        #             ρ₂=sys.ρ[i, j, 2],
        #             tait=tait_C₅H₁₂,
        #             igas=ideal_gas) where T<:AbstractFloat
        #

        ## Type unstable
        sys.P[i, j], sys.s[i, j, 1] =
                   find_pressure(p=p,
                   root_finder=newton_raphson,
                   ρ₁=sys.ρ[i, j, 1],
                   ρ₂=sys.ρ[i, j, 2],
                   tait=tait_C₅H₁₂,
                   igas=ideal_gas)

        sys.s[i, j, 2] = 1 - sys.s[i, j, 1]
    end
end

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

## -------------------- Saturation ---------------------- ##
##
## We derive the saturation from the molar composition
## ψ = ν₁ / ν₂ and the densities ρ̂₁ and ρ̂₂.

function saturation(ρ̂₁::T, ρ̂₂::T, ψ::T,
        M₁::T, M₂::T) where T<:AbstractFloat
    return (ρ̂₁ / ρ̂₂ * M₂ * (1 - ψ) / M₁ / ψ + 1)^(-1)
end

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


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

@testset "Dome pressures" begin
    #testing with R-value of zero to ignore the adjustment to PTRO
    @test domepressure_downhole(600-14.65, 0, 180) ≈ 750 atol = 5 #function takes psig but test values are psia
    @test domepressure_downhole(727-14.65, 0, 120) ≈ 820 atol = 5
end

@testset "Thornhill-Craver" begin
    @test ThornhillCraver_gaspassage_simplified(900, 1100, 140, 16) ≈ 1127 atol=1 #using psig inputs
    @test ThornhillCraver_gaspassage_simplified(1200, 1100, 140, 16) == 0

    seats = [7*4, 32, 5*8, 3*16]
    Hernandez_results = [3.15, 4.11, 6.36, 9.39] .* 1000
    results = [ThornhillCraver_gaspassage(420 - pressure_atmospheric, 850 - pressure_atmospheric, 150, s, 0.7) for s in seats]
    @test all(abs.(Hernandez_results .- results) .<= Hernandez_results .* 0.02) #2% tolerance to account for changing C_ds that aren't easily available (test cases use Winkler C_ds)

    @test ThornhillCraver_gaspassage(1200, 1100, 140, 16, 0.7) == 0
end

@testset "Valve table" begin

#EHU 256H example using Weatherford method

MDs = [0,1813, 2375, 2885, 3395]
TVDs = [0,1800, 2350, 2850, 3350]
incs = [0,0,0,0,0]
id = 2.441

well = Wellbore(MDs, incs, TVDs, id)
valves = GasliftValves([1813,2375,2885,3395], [1005,990,975,960], [0.073,0.073,0.073,0.073], [16,16,16,16])

tubing_pressures = [150,837,850,840,831]
casing_pressures = 1070 .+ [0,53,70,85,100]
temps = [135,145,148,151,153]

vdata, inj_depth = valve_calcs(valves = valves, well = well, sg_gas = 0.72, tubing_pressures = tubing_pressures, casing_pressures = casing_pressures, tubing_temps = temps, casing_temps = temps)

valve_table(vdata, inj_depth) #implicit test

results = vdata[1:4, [5,13,12,4]] #PSC, PVC, PVO, PSO

expected_results =
[1050. 1103 1124 1071;
 1023 1092 1111 1042;
 996  1080 1099 1015;
 968  1068 1087 987]

@test all(abs.(expected_results .- results) .< (expected_results .* 0.01)) #1% tolerance due to using TCFs versus PVT-based dome correction, as well as rounding errors

end #testset for valve table


"""
    theta(z, R, char=[], derivs=[], derivs_t=[])

Compute the Riemann theta function of the matrix in R at z, with optional characteristics or derivatives.

The input derivatives must be unit vectors, and the input z must be of the form ``z = a + τb, b∈[0,1)^g, a∈ \\mathbb{R}^g``.

# Arguments
- `z::Array{<:Number}`: Array of size g, each entry is a complex number.
- `R::RiemannMatrix`: RiemannMatrix type containing matrix.
- `char::Array{}=[]`: 2-element array containing 2 arrays of size g, where entries are either 0 or 1.
- `derivs::Array{}=[]`: N-element array specifying the derivative to the N-th order with respect to z, where each element is an array representing a unit vector in the direction of the derivative.
- `derivs_t::Array{}=[]`: M-element array specifying the derivative to the M-th order with respect to τ, where each element is an array of size 2 whose entries are the index in τ.

# Examples
```julia
julia> theta([0,0], RiemannMatrix([1+im -1; -1 1+im]), char=[[1,0],[1,1]], derivs=[[1,0]], derivs_t=[[2,1]])
```
"""
function theta(z::Array{<:Number}, R::RiemannMatrix; char::Array{}=[], derivs::Array{}=[], derivs_t::Array{}=[])
    # compute theta function with characteristics
    shift_x = zeros(R.g); # shift z by characteristic
    shift_n = zeros(R.g); # shift points in ellipsoid by characteristic
    if length(char) > 0
        shift_x = 0.5*char[2];
        shift_n = 0.5*char[1];
    end
    # compute derivatives with respect to τ by converting to derivatives with respect to z using the heat equation
    deriv_t_factor = 1; # multiply result by this factor
    for d in derivs_t
        if d[1] == d[2]
            deriv_t_factor *= 1/(4*π*im);
        else
            deriv_t_factor *= 1/(2*π*im);
        end
        deriv = [zeros(R.g), zeros(R.g)];
        deriv[1][d[1]] = 1;
        deriv[2][d[2]] = 1;
        derivs = vcat(derivs, deriv);
    end
    # compute theta function
    x = real(z) + shift_x; # shift z by characteristic
    y = convert(Array{Float64}, imag(z));
    y0 = (R.g > 1 ?  inv(R.Y)*y : inv.(R.Y).*y); # Y⁻¹y
    exp_part = exp((π*transpose(y)*y0)[1]); # exponential part of theta function
    osc_part = oscillatory_part(R, x, y0, shift_n, derivs); # oscillatory part of theta function
    return deriv_t_factor*exp_part*osc_part;
end

"""
    oscillatory_part(R, x, y0, shift_n, derivs=[])

Compute the oscillatory part of the Riemann theta function. Helper function to [`theta`](@ref).
"""
function oscillatory_part(R::RiemannMatrix, x::Array{<:Real}, y0::Array{<:Real}, shift_n::Array{<:Real}, derivs::Array{}=[])
    nderivs = size(derivs)[1];
    E = [n + shift_n for n in R.ellipsoid[nderivs + 1]]; # shift points in ellipsoid for computing theta function with characteristics
    w = round.(y0);
    w0 = y0 - w;
    s = (2*π*im)^(nderivs) * sum([
(nderivs > 0 ? prod(transpose(d)*(p-w) for d in derivs) : 1) * exp(2*π*im*((0.5*(transpose(p-w)*R.X*(p-w))[1] + transpose(p-w)*x)[1]) - π*(transpose(p+w0)*R.Y*(p+w0))[1]) for p in E]); # compute sum over points in ellipsoid
    return s;
end


using GLM
glm(:(y ~ 1), DataFrame(y = float64(randbool(10))), Binomial())

n = 2_500_000; srand(1234321)
df2 = DataFrame(x1 = rand(Normal(), n),
                x2 = rand(Exponential(), n),
                ss = pool(rand(DiscreteUniform(50), n)));
beta = unshift!(rand(Normal(),52), 0.5);
eta = ModelMatrix(ModelFrame(Formula(:(~ (x1 + x2 + ss))), df2)).m * beta;
mu = linkinv!(LogitLink(), similar(eta), eta);
df2["y"] = float64(rand(n) .< mu); df2["eta"] = eta; df2["mu"] = mu;
head(df2)
gm6 = glm(:(y ~ x1 + x2 + ss), df2, Binomial(), LogitLink())
@time glm(:(y ~ x1 + x2 + ss), df2, Binomial(), LogitLink());


using ClutteredEnvPathOpt
using LinearAlgebra
using Test
using Pipe
using Plots
using JuMP
using Gurobi

################################# solve_deits ########################################
# Create obstacles
num_obs = 3;
seed = 10;
obstacles = ClutteredEnvPathOpt.gen_field(num_obs, seed = seed)
plot();
ClutteredEnvPathOpt.plot_field(obstacles)
# points = ClutteredEnvPathOpt.find_points(obstacles)
# ClutteredEnvPathOpt.plot_points(points)
display(plot!())

# Set parameters
N = 30  # number of steps
f1 = [0.0, 0.1, 0.0]  # initial footstep pose 1
f2 = [0.0, 0.0, 0.0]  # initial footstep pose 2
goal = [1.0, 1.0, 0.0]  # goal pose
Q_g = 10 * Matrix{Float64}(I, 3, 3)  # weight between final footstep and goal pose
Q_r = [1 0 0; 0 1 0; 0 0 0] # Matrix{Float64}(I, 3, 3)  # weight between footsteps
q_t = -.05  # weight for trimming unused steps
# TODO: Assure second footstep is within reachability given first

# Optional named arguments
d1 = 0.2
d2 = 0.2
p1 = [0. 0.07]
p2 = [0, -0.27]
delta_x_y_max = 0.1  # max stride norm in space
delta_θ_max = pi/4  # max difference in θ
L = 5  # number of pieces in p.w.l approx. of sine/cosine

# Compute optimal path
x1, y1, θ1, t1, stats1 = solve_deits(obstacles, N, f1, f2, goal, Q_g, Q_r, q_t, method="merged")
x = copy(x1); y = copy(y1); θ = copy(θ1); t = copy(t1);
term_status, r_solve_time, rel_gap, simplex_iters, r_node_count, num_vertices, merged_size, full_size, num_free_faces, num_free_face_ineq, method_used = stats1

x2, y2, θ2, t2, stats2 = solve_deits(obstacles, N, f1, f2, goal, Q_g, Q_r, q_t, method="full")
term_status2, r_solve_time2, rel_gap2, simplex_iters2, r_node_count2, num_vertices2, merged_size2, full_size2, num_free_faces2, num_free_face_ineq2, method_used2 = stats2

x3, y3, θ3, t3, stats3 = solve_deits(obstacles, N, f1, f2, goal, Q_g, Q_r, q_t, method="bigM")
term_status3, r_solve_time3, rel_gap3, simplex_iters3, r_node_count3, num_vertices3, merged_size3, full_size3, num_free_faces3, num_free_face_ineq3, method_used3 = stats3

seed_range = 1:3
num_obs_range = 1:4
time_diffs = []
for seed in seed_range
    for num_obs in num_obs_range
        println("\nOn test Seed = $seed, Num_Obs = $num_obs")
        obstacles = ClutteredEnvPathOpt.gen_field(num_obs, seed = seed)
        _,_,_,_,stats = solve_deits(obstacles, N, f1, f2, goal, Q_g, Q_r, q_t, method = "merged")
        _,_,_,_,stats2 = solve_deits(obstacles, N, f1, f2, goal, Q_g, Q_r, q_t, method="full")
        _,_,_,_,stats3 = solve_deits(obstacles, N, f1, f2, goal, Q_g, Q_r, q_t, method="bigM")
        term_status, solve_time, rel_gap, simplex_iters, node_count, num_vertices, merged_size, full_size, num_free_faces, num_free_face_ineq, method_used = stats
        term_status2, solve_time2, rel_gap2, simplex_iters2, node_count2, num_vertices2, merged_size2, full_size2, num_free_faces2, num_free_face_ineq2, method_used2 = stats2
        term_status3, solve_time3, rel_gap3, simplex_iters3, node_count3, num_vertices3, merged_size3, full_size3, num_free_faces3, num_free_face_ineq3, method_used3 = stats3
        # Merged biclique times vs bigM
        t_diff = solve_time3 - solve_time
        push!(time_diffs, t_diff)
    end
end

# Trim excess steps
num_to_trim = length(filter(tj -> tj > 0.5, t[3:end]))
if num_to_trim % 2 == 0
    x = vcat(x[1:2], x[num_to_trim + 3 : end]);
    y = vcat(y[1:2], y[num_to_trim + 3 : end]);
    θ = vcat(θ[1:2], θ[num_to_trim + 3 : end]);
else
    x = vcat(x[1], x[num_to_trim + 3 : end]);
    y = vcat(y[1], y[num_to_trim + 3 : end]);
    θ = vcat(θ[1], θ[num_to_trim + 3 : end]);
end

num_to_trim2 = length(filter(tj -> tj > 0.5, t2[3:end]))
if num_to_trim2 % 2 == 0
    x2 = vcat(x2[1:2], x2[num_to_trim2 + 3 : end]);
    y2 = vcat(y2[1:2], y2[num_to_trim2 + 3 : end]);
    θ2 = vcat(θ2[1:2], θ2[num_to_trim2 + 3 : end]);
else
    x2 = vcat(x2[1], x2[num_to_trim2 + 3 : end]);
    y2 = vcat(y2[1], y2[num_to_trim2 + 3 : end]);
    θ2 = vcat(θ2[1], θ2[num_to_trim2 + 3 : end]);
end

num_to_trim3 = length(filter(tj -> tj > 0.5, t3[3:end]))
if num_to_trim3 % 2 == 0
    x3 = vcat(x3[1:2], x3[num_to_trim3 + 3 : end]);
    y3 = vcat(y3[1:2], y3[num_to_trim3 + 3 : end]);
    θ3 = vcat(θ3[1:2], θ3[num_to_trim3 + 3 : end]);
else
    x3 = vcat(x3[1], x3[num_to_trim3 + 3 : end]);
    y3 = vcat(y3[1], y3[num_to_trim3 + 3 : end]);
    θ3 = vcat(θ3[1], θ3[num_to_trim3 + 3 : end]);
end

# Plot footstep plan
plot_steps(obstacles, x, y, θ)
plot_steps(obstacles, x2, y2, θ2)
plot_steps(obstacles, x3, y3, θ3)
png("Path Seed $seed Num Obs $num_obs")

# Plot intersections of circles
plot_circles(x, y, θ)
plot_circles(x2, y2, θ2)
plot_circles(x3, y3, θ3)

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

num_obs = 4;
seed = 3;
obstacles, points, g, obstacle_faces, free_faces = ClutteredEnvPathOpt.plot_new(num_obs, "Obstacles Seed $seed Num Obs $num_obs", seed = seed)
skeleton = LabeledGraph(g)
all_faces = union(obstacle_faces, free_faces)

# Plot obstacles
plot(title="Obstacles");
ClutteredEnvPathOpt.plot_field(obstacles)
# points = ClutteredEnvPathOpt.find_points(obstacles)
ClutteredEnvPathOpt.plot_points(points, vertices=skeleton.labels)
# ClutteredEnvPathOpt.plot_lines(obstacles)
# ClutteredEnvPathOpt.plot_borders()
# ClutteredEnvPathOpt.plot_intersections(obstacles);
png("Obstacles Seed $seed Num Obs $num_obs")
display(plot!())

# Plot faces
ClutteredEnvPathOpt.plot_faces(obstacle_faces, points, plot_name = "Obstacle Faces", col = "dodgerblue")
png("Obstacle Faces Seed $seed Num Obs $num_obs")

ClutteredEnvPathOpt.plot_faces(free_faces, points, plot_name = "Free Faces")
png("Free Faces Seed $seed Num Obs $num_obs")

ClutteredEnvPathOpt.plot_faces(obstacle_faces, points, plot_name = "All Faces", col = "red")
ClutteredEnvPathOpt.plot_faces(free_faces, points, plot_name = "All Faces", col = "green", new_plot = false)
png("All Faces Seed $seed Num Obs $num_obs")

# Plot faces individually
ClutteredEnvPathOpt.plot_faces(obstacle_faces, points, col = "dodgerblue", individually = true)
ClutteredEnvPathOpt.plot_faces(free_faces, points, individually = true)

# Plot edges
ClutteredEnvPathOpt.plot_edges(skeleton, points, plot_name = "Skeleton", vertices=skeleton.labels)
png("Skeleton Seed $seed Num Obs $num_obs")

feg_S = ClutteredEnvPathOpt._find_finite_element_graph(skeleton, ClutteredEnvPathOpt._find_face_pairs(free_faces))
ClutteredEnvPathOpt.plot_edges(feg_S, points, plot_name = "Finite Element Graph of S",col="black", vertices=feg_S.labels)
png("FEG of S Seed $seed Num Obs $num_obs")

# feg_obs = ClutteredEnvPathOpt._find_finite_element_graph(skeleton, ClutteredEnvPathOpt._find_face_pairs(obstacle_faces))
# ClutteredEnvPathOpt.plot_edges(feg_obs, points, plot_name = "Finite Element Graph Obstacle Faces")
# png("FEG Obstacles Seed $seed Num Obs $num_obs")

# feg_all = ClutteredEnvPathOpt._find_finite_element_graph(skeleton, ClutteredEnvPathOpt._find_face_pairs(all_faces))
# ClutteredEnvPathOpt.plot_edges(feg_all, points, plot_name = "Finite Element Graph")
# png("FEG All Seed $seed Num Obs $num_obs")

# # Duplicate tests
# obstacles = ClutteredEnvPathOpt.gen_field(num_obs, seed = seed)
# points, mapped, inside_quant = ClutteredEnvPathOpt.find_intersections(obstacles)
# neighbors = ClutteredEnvPathOpt._find_neighbors(points, mapped)
# points = ClutteredEnvPathOpt.find_points(obstacles)

# dup, dup_ind = face_duplicates(obstacle_faces)
# dup, dup_ind = face_duplicates(free_faces)
# dup, dup_ind = face_duplicates(all_faces)
# plot_intersections_indiv(points)
# dup, dup_ind = point_duplicates(points)
# dup, dup_ind = points_in_hs_duplicates(mapped)
# dup, dup_ind = hs_duplicates(mapped)
# points_in_mapped_ordered(obstacles, mapped)
# plot_neighbors(obstacles, neighbors, points)
# list_neighbors(neighbors, points)
# suspect = suspect_neighbors(neighbors)

# cover3 = ClutteredEnvPathOpt._find_biclique_cover_visual(skeleton, free_faces, points)
# tree = ClutteredEnvPathOpt.find_biclique_cover_as_tree(skeleton, free_faces)

# Free Faces Cover
cover = ClutteredEnvPathOpt.find_biclique_cover(skeleton, free_faces)
(valid_cover, size_diff, missing_edges, missing_edges_lg) = ClutteredEnvPathOpt._is_valid_biclique_cover_diff(ClutteredEnvPathOpt._find_finite_element_graph(skeleton, ClutteredEnvPathOpt._find_face_pairs(free_faces)), cover)

feg_free = ClutteredEnvPathOpt._find_finite_element_graph(skeleton, ClutteredEnvPathOpt._find_face_pairs(free_faces))
merged_cover = ClutteredEnvPathOpt.biclique_merger(cover, feg_free)
(valid_cover, size_diff, missing_edges, missing_edges_lg) = ClutteredEnvPathOpt._is_valid_biclique_cover_diff(ClutteredEnvPathOpt._find_finite_element_graph(skeleton, ClutteredEnvPathOpt._find_face_pairs(free_faces)), merged_cover)

ClutteredEnvPathOpt.plot_edges(missing_edges_lg, points, plot_name = "Missing Edges",vertices=missing_edges_lg.labels)
png("Missing Edges Seed $seed Num Obs $num_obs Free")

# file_name = "Biclique Cover Debug Output Seed $seed Num Obs $num_obs Free Unedited.txt"
# f = open(file_name, "w")
# write(f, "Plot 1\n")
# flush(f)
# close(f)
# cover2 = ClutteredEnvPathOpt.find_biclique_cover_debug(skeleton, free_faces, points, obstacles, file_name)

# To produce individual FEGs/Skeletons in the recursion process
(C,A,B), skeleton_ac, faces_ac, skeleton_bc, faces_bc = ClutteredEnvPathOpt.find_biclique_cover_one_iter(skeleton, free_faces)
feg_S_ac = ClutteredEnvPathOpt._find_finite_element_graph(skeleton_ac, ClutteredEnvPathOpt._find_face_pairs(faces_ac))
ClutteredEnvPathOpt.plot_edges(feg_S_ac, points, plot_name = "Finite Element Graph of S_ac",vertices=feg_S_ac.labels,col="black")
# plot!(title="",axis=([], false))
png("FEG of S_ac Seed $seed Num Obs $num_obs")
feg_S_bc = ClutteredEnvPathOpt._find_finite_element_graph(skeleton_bc, ClutteredEnvPathOpt._find_face_pairs(faces_bc))
ClutteredEnvPathOpt.plot_edges(feg_S_bc, points, plot_name = "Finite Element Graph of S_bc", vertices=feg_S_bc.labels, col="black")
# plot!(title="",axis=([], false))
png("FEG of S_bc Seed $seed Num Obs $num_obs")

# Biclique Cover Validity Tests-----------------------------------------------------------------------
seed_range = 1:100
num_obs_range = 1:4
# fails, tuples = biclique_cover_validity_tests(seed_range, num_obs_range, faces = "all")
fails, tuples = biclique_cover_validity_tests(seed_range, num_obs_range, faces = "free", with_plots=true)

percent_decreases, tuples, c_sizes, mc_sizes = biclique_cover_merger_stats(seed_range, num_obs_range)
percent_decreases_hp, tuples_hp, c_sizes_hp, mc_sizes_hp = biclique_cover_merger_stats(seed_range, num_obs_range, file_name="Biclique Cover Merging Stats HP Partition.txt", partition="hp")
size_diff_cover = c_sizes_hp - c_sizes
size_diff_merged_cover = mc_sizes_hp - mc_sizes

# TODO: WRITE STATS TO FILE

# Biclique validity tester function
function biclique_cover_validity_tests(seed_range, num_obs_range; faces::String="free", with_plots::Bool=false)
    fails = 0
    tuples = []
    for seed = seed_range
        for num_obs = num_obs_range
            println("On test Seed = $seed, Num_Obs = $num_obs")
            obstacles, points, g, obstacle_faces, free_faces = ClutteredEnvPathOpt.plot_new(num_obs,"Obstacles Seed $seed Num Obs $num_obs",seed=seed); # optional seed argument, default is 11
            skeleton = LabeledGraph(g)
            all_faces = union(obstacle_faces, free_faces)
            if faces == "all"
                cover = ClutteredEnvPathOpt.find_biclique_cover(skeleton, all_faces);
                feg = ClutteredEnvPathOpt._find_finite_element_graph(skeleton, ClutteredEnvPathOpt._find_face_pairs(all_faces));
                merged_cover = ClutteredEnvPathOpt.biclique_merger(cover, feg)
                # (valid_cov, size_diff, miss_edges, miss_edges_lg) = ClutteredEnvPathOpt._is_valid_biclique_cover_diff(feg, cover);
                (valid_cov, size_diff, miss_edges, miss_edges_lg) = ClutteredEnvPathOpt._is_valid_biclique_cover_diff(feg, merged_cover);
                println("Difference: $(length(cover)-length(merged_cover)) | Merged: $(length(merged_cover)) | Original: $(length(cover))")
            else
                cover = ClutteredEnvPathOpt.find_biclique_cover(skeleton, free_faces);
                feg = ClutteredEnvPathOpt._find_finite_element_graph(skeleton, ClutteredEnvPathOpt._find_face_pairs(free_faces));
                merged_cover = ClutteredEnvPathOpt.biclique_merger(cover, feg)
                # (valid_cov, size_diff, miss_edges, miss_edges_lg) = ClutteredEnvPathOpt._is_valid_biclique_cover_diff(feg, cover);
                (valid_cov, size_diff, miss_edges, miss_edges_lg) = ClutteredEnvPathOpt._is_valid_biclique_cover_diff(feg, merged_cover);
                println("Difference: $(length(cover)-length(merged_cover)) | Merged: $(length(merged_cover)) | Original: $(length(cover))")
            end
            if size_diff != 0
                fails += 1
                push!(tuples, (seed, num_obs))
                println("Test Failed")

                if with_plots
                    plot(title="Obstacles");
                    ClutteredEnvPathOpt.plot_field(obstacles)
                    # points = ClutteredEnvPathOpt.find_points(obstacles)
                    ClutteredEnvPathOpt.plot_points(points)
                    # ClutteredEnvPathOpt.plot_lines(obstacles)
                    # ClutteredEnvPathOpt.plot_borders()
                    # ClutteredEnvPathOpt.plot_intersections(obstacles);
                    display(plot!())
                    png("Obstacles Seed $seed Num Obs $num_obs")

                    ClutteredEnvPathOpt.plot_faces(free_faces, points, plot_name = "Free Faces")
                    png("Free Faces Seed $seed Num Obs $num_obs")

                    ClutteredEnvPathOpt.plot_edges(skeleton, points, plot_name = "Skeleton")
                    png("Skeleton Seed $seed Num Obs $num_obs")

                    if size_diff > 0
                        ClutteredEnvPathOpt.plot_edges(miss_edges_lg, points, plot_name = "Missing Edges")
                        if faces == "all"
                            png("Missing Edges All Faces Used Seed $seed Num Obs $num_obs")
                        else
                            png("Missing Edges Free Faces Used Seed $seed Num Obs $num_obs")
                        end
                    end
                end
            end
        end
    end
    return (fails, tuples)
end

# Biclique cover merger function for stats
function biclique_cover_merger_stats(seed_range, num_obs_range; file_name="Biclique Cover Merging Stats.txt", partition="CDT")
    # file_name = "Biclique Cover Merging Stats.txt" #Seed Range = $seed_range, Num Obs Range = $num_obs_range.txt"
    f = open(file_name, "w")
    write(f, "Seed Range = $seed_range, Num Obs Range = $num_obs_range\n")
    flush(f)
    close(f)

    f = open(file_name, "a")
    write(f, "Seed\tNum_obs\tFull_cover\tMerged_cover\tDecrease\tPercent_decrease\n")
    fails = 0
    tuples = []
    percent_vec = []
    cover_vec = []
    merged_cover_vec = []
    for seed = seed_range
        for num_obs = num_obs_range
            # f = open(file_name, "a")
            # write(f, "\nTest Seed = $seed, Num_Obs = $num_obs\n")
            write(f, "$seed\t$num_obs")
            # flush(f)
            # close(f)

            # println("On test Seed = $seed, Num_Obs = $num_obs")
            obstacles, points, g, obstacle_faces, free_faces = ClutteredEnvPathOpt.plot_new(num_obs,"Obstacles Seed $seed Num Obs $num_obs",seed=seed;partition=partition); # optional seed argument, default is 11
            skeleton = LabeledGraph(g)
            all_faces = union(obstacle_faces, free_faces)

            cover = ClutteredEnvPathOpt.find_biclique_cover(skeleton, free_faces);
            feg = ClutteredEnvPathOpt._find_finite_element_graph(skeleton, ClutteredEnvPathOpt._find_face_pairs(free_faces));
            merged_cover = ClutteredEnvPathOpt.biclique_merger(cover, feg)
            (valid_cover, size_diff, miss_edges, miss_edges_lg) = ClutteredEnvPathOpt._is_valid_biclique_cover_diff(feg, cover);
            (valid_cover_merged, size_diff, miss_edges, miss_edges_lg) = ClutteredEnvPathOpt._is_valid_biclique_cover_diff(feg, merged_cover);
            # println("Difference: $(length(cover)-length(merged_cover)) | Merged: $(length(merged_cover)) | Original: $(length(cover))")

            # if valid_cover && valid_cover_merged
            #     write(f,"Difference: $(length(cover)-length(merged_cover)) | Merged: $(length(merged_cover)) | Original: $(length(cover))")
            #     percent = round( ((length(cover)-length(merged_cover)) / length(cover)) * 100, digits = 2)
            #     write(f," | Percent Decrease: $percent\n")
            #     push!(percent_vec, percent)
            # else
            #     fails += 1
            #     push!(tuples, (seed, num_obs))
            #     if !valid_cover_merged
            #         write(f,"Merged cover invalid.\n")
            #     else
            #         write(f,"Original cover invalid.\n")
            #     end
            # end
            if valid_cover && valid_cover_merged
                percent = round( ((length(cover)-length(merged_cover)) / length(cover)) * 100, digits = 2)
                write(f,"\t$(length(cover))\t$(length(merged_cover))\t$(length(cover)-length(merged_cover))\t$percent\n")
                push!(percent_vec, percent)
                push!(cover_vec, length(cover))
                push!(merged_cover_vec, length(merged_cover))
            elseif valid_cover
                percent = round( ((length(cover)-0) / length(cover)) * 100, digits = 2)
                write(f,"\t$(length(cover))\t0\t$(length(cover)-0)\t$percent\n")
                push!(percent_vec, percent)
                push!(cover_vec, -1)
                push!(merged_cover_vec, -1)

                fails += 1
                push!(tuples, (seed, num_obs))
                # if !valid_cover_merged
                #     write(f,"Merged cover invalid.\n")
                # else
                #     write(f,"Original cover invalid.\n")
                # end
            end
        end
    end

    # write(f, "\nTotal Fails: $fails.\nTuples: $tuples")
    flush(f)
    close(f)

    return percent_vec, tuples, cover_vec, merged_cover_vec
end

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

# method <- "merged" for the compact biclique cover, "full" for the original biclique cover, "bigM" for big-M constraints
function solve_time_stats(seed_range, num_obs_range; method="merged", file_name="Solve Time Stats.txt")
    # file_name = "Solve Time Stats.txt" #Seed Range = $seed_range, Num Obs Range = $num_obs_range.txt"
    # f = open(file_name, "a")
    # write(f, "Method = $method, Seed Range = $seed_range, Num Obs Range = $num_obs_range\n")
    # flush(f)
    # close(f)

    # Set parameters
    N = 30  # number of steps
    f1 = [0.0, 0.1, 0.0]  # initial footstep pose 1
    f2 = [0.0, 0.0, 0.0]  # initial footstep pose 2
    goal = [1, 1, 0]  # goal pose
    Q_g = 10*Matrix{Float64}(I, 3, 3)  # weight between final footstep and goal pose
    Q_r = [1 0 0; 0 1 0; 0 0 0] # Matrix{Float64}(I, 3, 3)  # weight between footsteps
    q_t = -.05  # weight for trimming unused steps
    # Optional named arguments
    # d1 = 0.2 <- radius of reference foot circle
    # d2 = 0.2 <- radius of moving foot circle
    # p1 = [0, 0.07] <- center of reference foot circle
    # p2 = [0, -0.27] <- center of moving foot circle
    # delta_x_y_max = 0.10  # max stride norm in space
    # delta_θ_max = pi/4  # max difference in θ
    # L = 5  # number of pieces in p.w.l approx. of sine/cosine

    f = open(file_name, "a")
    # write(f, "Seed\tNum_obs\tTime\n")
    write(f, "Seed\tNum_obs\tSolve_time\tRel_gap\tSimplex_iterations\tNodes_explored\tNum_vertices\tMerged_size\tFull_size\tNum_free_faces\tNum_free_face_inequalities\n")
    times = zeros(length(seed_range) * length(num_obs_range))
    i = 1
    for seed = seed_range
        for num_obs = num_obs_range
            # f = open(file_name, "a")
            # write(f, "Seed = $seed, Num_Obs = $num_obs\n")
            # flush(f)
            # close(f)
            # println("On test Seed = $seed, Num_Obs = $num_obs")

            # Create obstacles
            obstacles = ClutteredEnvPathOpt.gen_field(num_obs, seed = seed)

            x, y, θ, t, stats = solve_deits(obstacles, N, f1, f2, goal, Q_g, Q_r, q_t, method=method)
            term_status, r_solve_time, rel_gap, simplex_iters, r_node_count, num_vertices, merged_size, full_size, num_free_faces, num_free_face_ineq, method_used = stats

            if method_used == method
                times[i] = r_solve_time # time
            else
                times[i] = -1
            end
            write(f, "$seed\t$num_obs\t$(times[i])\t$rel_gap\t$simplex_iters\t$r_node_count\t$num_vertices\t$merged_size\t$full_size\t$num_free_faces\t$num_free_face_ineq\n")
            flush(f)

            i += 1

            # # Trim excess steps
            # num_to_trim = length(filter(tj -> tj > 0.5, t[3:end]))
            # if num_to_trim % 2 == 0
            #     x = vcat(x[1:2], x[num_to_trim + 3 : end]);
            #     y = vcat(y[1:2], y[num_to_trim + 3 : end]);
            #     θ = vcat(θ[1:2], θ[num_to_trim + 3 : end]);
            # else
            #     x = vcat(x[1], x[num_to_trim + 3 : end]);
            #     y = vcat(y[1], y[num_to_trim + 3 : end]);
            #     θ = vcat(θ[1], θ[num_to_trim + 3 : end]);
            # end
            # plot_steps(obstacles, x, y, θ)
            # plot_circles(x, y, θ, p1=p1, p2=p2)

        end
    end

    # write(f, "\nEND")
    # flush(f)
    close(f)

    return times

end

seed_range = 21:30 # done 1:20
num_obs_range = 1:4
methods = ("merged", "full", "bigM")
times = Dict()
# method = "merged"
for method in methods
    file_name = "Solve Time Stats Method $method.txt"
    f = open(file_name, "a")   # write or append appropriately
    write(f, "Seed Range = $seed_range, Num Obs Range = $num_obs_range\nMethod = $method\n")
    flush(f)
    close(f)
    time = solve_time_stats(seed_range, num_obs_range, method=method, file_name=file_name)
    times[method] = time
end

file_name = "Solve Time Stats Comparison.txt" #Seed Range = $seed_range, Num Obs Range = $num_obs_range.txt"
f = open(file_name, "a")    # write or append appropriately
times_m_b = times["bigM"] .- times["merged"]
times_m_f = times["full"] .- times["merged"]
times_f_b = times["bigM"] .- times["full"]
write(f, "Seed Range = $seed_range, Num Obs Range = $num_obs_range\n")
write(f, "Merged_vs_BigM\tMerged_vs_Full\tFull_vs_BigM\n")
for i = 1:length(times_m_b)
    write(f, "$(times_m_b[i])\t$(times_m_f[i])\t$(times_f_b[i])\n")
end
flush(f)
close(f)

import DelimitedFiles
data_merged = DelimitedFiles.readdlm("Solve Time Stats Method merged.txt", '\t')
matrix_merged = Matrix{Float64}(data_merged[4:end,:])
data_bigM = DelimitedFiles.readdlm("Solve Time Stats Method bigM.txt", '\t')
matrix_bigM = Matrix{Float64}(data_bigM[4:end,:])
header = Vector{String}(data_merged[3,:])

file_name = "Solve Stats Comparison.txt" #Seed Range = $seed_range, Num Obs Range = $num_obs_range.txt"
f = open(file_name, "w")    # write or append appropriately
write(f, "S_T_m\tS_T_b\tS_T_diff\tS_I_m\tS_I_b\tS_I_diff\tN_C_m\tN_C_b\tN_C_diff\n")
S_T_m = matrix_merged[:,3]
S_T_b = matrix_bigM[:,3]
S_T_diff = S_T_b .- S_T_m
S_I_m = matrix_merged[:,5]
S_I_b = matrix_bigM[:,5]
S_I_diff = S_I_b .- S_I_m
N_C_m = matrix_merged[:,6]
N_C_b = matrix_bigM[:,6]
N_C_diff = N_C_b .- N_C_m
for i = 1:length(S_T_m)
    write(f, "$(S_T_m[i])\t$(S_T_b[i])\t$(S_T_diff[i])\t")
    write(f, "$(S_I_m[i])\t$(S_I_b[i])\t$(S_I_diff[i])\t")
    write(f, "$(N_C_m[i])\t$(N_C_b[i])\t$(N_C_diff[i])\n")
end
flush(f)
close(f)

data_summary = DelimitedFiles.readdlm("Solve Stats Comparison.txt", '\t')
matrix_summary = Matrix{Float64}(data_summary[2:end,:])
header_summary = Vector{String}(data_summary[1,:])

# Debugging Functions ---------------------------------------------------------------

# Check if free_faces has any duplicates---------------------------------------------
function face_duplicates(faces)
    dup = 0
    dup_ind = []
    col_faces = collect(faces)
    for i in 1:length(col_faces)
        for j = (i+1):length(col_faces)
            sort_i = sort(col_faces[i])
            sort_j = sort(col_faces[j])
            if sort_i == sort_j
                dup += 1
                push!(dup_ind, (i,j))
            end
        end
    end
    return dup, dup_ind
end

# Check if points has any duplicates---------------------------------------------
function point_duplicates(points)
    dup = 0
    dup_ind = []
    tol = Rational(1.e-13)
    for i in 1:(length(points))
        for j in (i+1):(length(points))
            #if (points[i].first == points[j].first) && (points[i].second == points[j].second)
            if abs(points[i].first - points[j].first) < tol && abs(points[i].second - points[j].second) < tol
                dup += 1
                push!(dup_ind, (i,j))
            end
        end
    end
    # for (i,j) in dup_ind
    #     println("Points $i and $j: \n ($(points[i])) \n ($(points[j]))")
    # end
    return dup, dup_ind
end

# Check if halfspaces in mapped have any duplicated points within them ---------
function points_in_hs_duplicates(mapped)
    dup = 0
    dup_ind = []
    for (l,hs) in enumerate(mapped)
        for i in 1:(length(hs))
            for j in (i+1):(length(hs))
                if (hs[i].first == hs[j].first) && (hs[i].second == hs[j].second)
                    dup += 1
                    push!(dup_ind, (l,i,j))
                end
            end
        end
    end
    return dup, dup_ind
end

# Check if any halfspaces in mapped are duplicates-----------------------------
function hs_duplicates(mapped)
    dup = 0
    dup_ind = []
    for i = 1:length(mapped)
        for j = (i+1):length(mapped)
            if mapped[i] == mapped[j]
                dup += 1
                push!(dup_ind, (i,j))
            end
        end
    end
    return dup, dup_ind
end

# Check to see if points are in order in mapped---------------------------------
function points_in_mapped_ordered(obstacles, mapped)
    for hs in mapped
        plot(legend=false,xlims=(-0.05,1.05),ylims=(-0.05,1.05));
        ClutteredEnvPathOpt.plot_field(obstacles)
        ClutteredEnvPathOpt.plot_lines(obstacles)
        for point in hs
            display(scatter!([point.first], [point.second]))
        end
    end
end

# Plot points individually----------------------------------------------------------
function plot_intersections_indiv(intersections)
    x = map(point -> point[1], intersections)
    y = map(point -> point[2], intersections)
    Plots.plot(legend=false,xlims=(-0.05,1.05),ylims=(-0.05,1.05))
    for i = 1:length(intersections)
        # Plots.plot(legend=false,xlims=(-0.05,1.05),ylims=(-0.05,1.05))
        display(Plots.scatter!([x[i]],[y[i]],title="Point $i:\n$(x[i])\n$(y[i])"))
    end
end

# Check if neighbors are correct----------------------------------------------------
function plot_neighbors(obstacles, neighbors, points)
    for (i,point) in enumerate(points)
        plot();
        ClutteredEnvPathOpt.plot_field(obstacles)
        ClutteredEnvPathOpt.plot_lines(obstacles)
        display(scatter!([point.first],[point.second],title="Point $i"))
        x_i_neighbors = map(node -> points[node].first, neighbors[i])
        y_i_neighbors = map(node -> points[node].second, neighbors[i])
        display(scatter!(x_i_neighbors,y_i_neighbors, title="Point $i. Neighbors $(neighbors[i])"))
        #println("Point $i has neighbors $(neighbors[i]).")
    end
end

function list_neighbors(neighbors, points)
    for i in 1:length(points)
        println("Point $i has neighbors $(neighbors[i]).")
    end
end

function duplicates(n)
    dup = 0
    for i in 1:length(n)
        for j in (i+1):length(n)
            if n[i] == n[j]
                dup += 1
            end
        end
    end
    return dup != 0
end

function suspect_neighbors(neighbors)
    suspect = []
    for n in 1:length(neighbors)
        if length(neighbors[n]) > 4 || duplicates(neighbors[n])
        #if duplicates(neighbors[n])
            push!(suspect, n)
        end
    end
    for point in suspect
        println("Point $point has neighbors $(neighbors[point])")
    end
    return suspect
end

# # Face finder ----------------------------------------------------------------------------
# t_seed = 7
# t_num_obs = 4
# obstacles = ClutteredEnvPathOpt.gen_field(t_num_obs, seed = t_seed) # optional seed
# construct_graph_debug(obstacles)


include(joinpath(@__DIR__, "../../src/FLS.jl"))

# Example: categories
catschema = @acset LabeledGraph_{Symbol} begin
  V = 7
  vlabel=[:O, :A, :cmp, :asc_l, :asc_r,
          :leftid, :rightid]
  E = 17
  elabel=[:src,:tgt, :refl,
      :c1, :c2, :c3,
      :l1, :l2, :r1, :r2, :asc,
      :lidv, :lida, :lidc,
      :ridv, :rida, :ridc,
      ]

  src = [2,2,1, 3,3,3, 4,4,5,5,4, 6,6,6, 7,7,7]
  tgt = [1,1,2, 2,2,2, 3,3,3,3,5, 1,2,3, 1,2,3]
end

"""Pair of arrows that match head to tail"""
cmpconed = @acset LabeledGraph_{Symbol} begin
  V = 3
  vlabel = [:A,:A,:O]
  E = 2
  elabel = [:src, :tgt]
  src = [1,2]
  tgt = [3,3]
end
cmpcone = Cone(cmpconed, :cmp, [1=>:c1, 2=>:c2])

"""(m₁⋅m₂)⋅m₃"""
asc_l_coned =  @acset LabeledGraph_{Symbol} begin
  V = 3
  E = 2
  vlabel = [:cmp, :cmp, :A]
  elabel = [:c3, :c1]
  src = [1,2]
  tgt = [3,3]
end

asc_l_cone = Cone(asc_l_coned,   :asc_l,  [1=>:l1, 2=>:l2])

"""m₁⋅(m₂⋅m₃)"""
asc_r_coned =  @acset LabeledGraph_{Symbol} begin
  V = 3
  E = 2
  vlabel = [:cmp, :cmp, :A]
  elabel = [:c2, :c3]
  src = [1,2]
  tgt = [3,3]
end
asc_r_cone = Cone(asc_r_coned,   :asc_r,   [1=>:r1, 2=>:r2])

"""(id(x)) ⋅ _ """
leftid_coned = @acset LabeledGraph_{Symbol} begin
  V = 3
  E = 2
  vlabel = [:O, :A, :cmp]
  elabel = [:c1, :refl]
  src = [3,1]
  tgt = [2,2]
end
leftid_cone = Cone(leftid_coned,  :leftid,  [1=>:lidv, 2=>:lida, 3=>:lidc])

rightid_coned = @acset LabeledGraph_{Symbol} begin
  V = 3
  E = 2
  vlabel = [:O, :A, :cmp]
  elabel = [:c2, :refl]
  src = [3,1]
  tgt = [2,2]
end
rightid_cone = Cone(rightid_coned, :rightid, [1=>:ridv, 2=>:rida, 3=>:ridc])

catcones = [cmpcone, asc_l_cone, asc_r_cone, leftid_cone, rightid_cone]

cateqs = [
  # reflexivity
  (:refl_src, [:refl, :src], Symbol[]),
  (:refl_tgt, [:refl, :tgt], Symbol[]),
  # Unitality
  (:unit_c2, [:lidc, :c2], [:lidc, :c3]),
  (:unit_c1, [:ridc, :c1], [:ridc, :c3]),
  # associativity
  (:l1c1, [:l1,:c1], [:asc, :r1, :c1]),
  (:l1c2, [:l1,:c2], [:asc, :r2, :c1]),
  (:l2c2, [:l2,:c2], [:asc, :r2, :c2]),
  (:l2c3, [:l2,:c3], [:asc, :r1, :c3])
]

catfls = FLS(:catt, catschema, catcones, cateqs);

cata2o1 = catfls.cset()
add_parts!(cata2o1, :O, 1)
add_parts!(cata2o1, :A, 2, src=[1,1], tgt=[1,1])
set_subpart!(cata2o1, 1, :refl, 1)
add_parts!(cata2o1, :cmp, 4, c1=[1,1,2,2], c2=[1,2,1,2], c3=[1,2,2,1])
add_parts!(cata2o1, :asc_r, 8, r1= [1,2,2,1,3,4,4,3],
                              # r1: 1st arg is l1.c1, last arg l1.c2⋅l2.c2
                            # l1c2:  1 1 2 2 1 1 2 2
                            # l2c2:  1 2 1 2 1 2 1 2
                            # l1c1:  1 1 1 1 2 2 2 2   <--
                     # l1.c2⋅l2.c2:  1 2 2 1 1 2 2 1   <--
                     r2= [1,2,3,4,1,2,3,4])
                     # r2: first arg is l1.c2, second arg is l2.c2

add_parts!(cata2o1, :asc_l, 8, l1= [1,1,2,2,3,3,4,4],
                              #l1c3: 1 1 2 2 3 3 4 4
                              #l1c3: 1 1 2 2 2 2 1 1 (must be 1st arg of l2)
                                l2= [1,2,3,4,3,4,1,2],
                                asc=[1,2,3,4,5,6,7,8]) # needs to be fixed.
add_parts!(cata2o1, :leftid, 2, lidv= [1,1], lida=[1,1], lidc=[1,2])
add_parts!(cata2o1, :rightid, 2, ridv= [1,1], rida=[1,1], ridc=[1,3]);
a2o1 = initrel(catfls, cata2o1);


module straight_duct_examples
using FinEtools
using FinEtoolsAcoustics
using UnicodePlots

function straight_duct_Q8_example()
    t0  =  time()

    rho = 1.21*phun("kg/m^3");# mass density
    c  = 343.0*phun("m/s");# sound speed
    bulk =  c^2*rho;
    omega =  54.5901*phun("rev/s")
    vn0 =  -1.0*phun("m/s")
    Lx = 10.0*phun("m");# length of the box, millimeters
    Ly = 1.0*phun("m"); # length of the box, millimeters
    n = 20;#number of elements along the length

    println("""

    Straight duct with anechoic termination.
    Example from Boundary element acoustics: Fundamentals and computer codes, TW Wu, page 44.
    Both real and imaginary components of the pressure should have amplitude of
    rho*c = $(rho*c).

    Two-dimensional model. Quadratic qquadrilateral mesh.
    """)

    fens,fes  =  Q8block(Lx,Ly,n,2); # Mesh
    bfes  =  meshboundary(fes)
    L0 = selectelem(fens,bfes,facing = true, direction = [-1.0 0.0])
    L10 = selectelem(fens,bfes,facing = true, direction = [+1.0 0.0])
    nLx = selectnode(fens,box = [0.0 Lx  0.0 0.0], inflate = Lx/1.0e5)

    geom  =  NodalField(fens.xyz)
    P  =  NodalField(zeros(ComplexF64,size(fens.xyz,1),1))

    numberdofs!(P)

    material = MatAcoustFluid(bulk,rho)
    femm  =  FEMMAcoust(IntegDomain(fes, GaussRule(2, 3)), material)

    S  =  acousticstiffness(femm, geom, P);
    C  =  acousticmass(femm, geom, P);

    E10femm  =  FEMMAcoustSurf(IntegDomain(subset(bfes,L10),GaussRule(1, 3)), material)
    D  =  acousticABC(E10femm, geom, P);

    E0femm  =  FEMMBase(IntegDomain(subset(bfes,L0), GaussRule(1,  3)))
    fi  =  ForceIntensity(-1.0im*omega*rho*vn0);
    F  =  distribloads(E0femm, geom, P, fi, 2);

    p = (-omega^2*S +omega*1.0im*D + C)\F
    scattersysvec!(P, p[:])

    println("Pressure amplitude bounds: ")
    println("  real $(minimum(real(P.values)))/$(maximum(real(P.values)))")
    println("  imag $(minimum(imag(P.values)))/$(maximum(imag(P.values)))")

    println("Total time elapsed  =  ",time() - t0,"s")

    File  =   "straight_duct_Q8.vtk"
    scalars = real(P.values);
    vtkexportmesh(File, connasarray(fes), geom.values, FinEtools.MeshExportModule.Q8;  scalars = [("Pressure", scalars)])
    @async run(`"paraview.exe" $File`)

    ix = sortperm(geom.values[nLx,1])
    # @pgf a = Axis({
    #           xlabel = "x",
    #           ylabel = L"$\mathrm{Re}P$, $\mathrm{Im}P$",
    #           title = "Steady-state pressure"
    #       },
    #       Plot(Table([:x => geom.values[nLx,1][ix], :y => real(P.values)[nLx][ix]])), LegendEntry("real"),
    #       Plot(Table([:x => geom.values[nLx,1][ix], :y => imag(P.values)[nLx][ix]])), LegendEntry("imag"))
    plt = lineplot(geom.values[nLx,1][ix], real(P.values)[nLx][ix], canvas = DotCanvas, title = "Steady-state pressure", name = "real", xlabel = "x", ylabel = "P")
    plt = lineplot!(plt, geom.values[nLx,1][ix], imag(P.values)[nLx][ix], name = "imag")
    display(plt)

    true

end # straight_duct_Q8_example

function straight_duct_T3_example()
    t0  =  time()

    rho = 1.21*phun("kg/m^3");# mass density
    c  = 343.0*phun("m/s");# sound speed
    bulk =  c^2*rho;
    omega =  54.5901*phun("rev/s")
    vn0 =  -1.0*phun("m/s")
    Lx = 10.0*phun("m");# length of the box, millimeters
    Ly = 1.0*phun("m"); # length of the box, millimeters
    n = 20;#number of elements along the length

    println("""

    Straight duct with anechoic termination.
    Example from Boundary element acoustics: Fundamentals and computer codes, TW Wu, page 44.
    Both real and imaginary components of the pressure should have amplitude of
    rho*c = $(rho*c).

    Triangle mesh.
    """)

    fens,fes  =  T3block(Lx,Ly,n,2); # Mesh
    bfes  =  meshboundary(fes)
    L0 = selectelem(fens,bfes,facing = true, direction = [-1.0 0.0])
    L10 = selectelem(fens,bfes,facing = true, direction = [+1.0 0.0])
    nLx = selectnode(fens,box = [0.0 Lx  0.0 0.0], inflate = Lx/1.0e5)

    geom  =  NodalField(fens.xyz)
    P  =  NodalField(zeros(ComplexF64,size(fens.xyz,1),1))

    numberdofs!(P)

    material = MatAcoustFluid(bulk,rho)
    femm  =  FEMMAcoust(IntegDomain(fes, TriRule(1)), material)

    S  =  acousticstiffness(femm, geom, P);
    C  =  acousticmass(femm, geom, P);

    E10femm  =  FEMMAcoustSurf(IntegDomain(subset(bfes,L10),GaussRule(1, 2)), material)
    D  =  acousticABC(E10femm, geom, P);

    E0femm  =  FEMMBase(IntegDomain(subset(bfes,L0), GaussRule(1,  2)))
    fi  =  ForceIntensity(-1.0im*omega*rho*vn0);
    F  =  distribloads(E0femm, geom, P, fi, 2);

    p = (-omega^2*S +omega*1.0im*D + C)\F
    scattersysvec!(P, p[:])

    println("Pressure amplitude bounds: ")
    println("  real $(minimum(real(P.values)))/$(maximum(real(P.values)))")
    println("  imag $(minimum(imag(P.values)))/$(maximum(imag(P.values)))")

    println("Total time elapsed  =  ",time() - t0,"s")

    File  =   "straight_duct.vtk"
    scalars = real(P.values);
    vtkexportmesh(File, connasarray(fes), geom.values, FinEtools.MeshExportModule.T3; scalars = [("Pressure", scalars)])
    @async run(`"paraview.exe" $File`)

    ix = sortperm(geom.values[nLx,1])
    # @pgf a = Axis({
    #           xlabel = "x",
    #           ylabel = L"$\mathrm{Re}P$, $\mathrm{Im}P$",
    #           title = "Steady-state pressure"
    #       },
    #       Plot(Table([:x => geom.values[nLx,1][ix], :y => real(P.values)[nLx][ix]])), LegendEntry("real"),
    #       Plot(Table([:x => geom.values[nLx,1][ix], :y => imag(P.values)[nLx][ix]])), LegendEntry("imag"))
    # display(a)
    plt = lineplot(geom.values[nLx,1][ix], real(P.values)[nLx][ix], canvas = DotCanvas, title = "Steady-state pressure", name = "real", xlabel = "x", ylabel = "P")
    plt = lineplot!(plt, geom.values[nLx,1][ix], imag(P.values)[nLx][ix], name = "imag")
    display(plt)

    true

end # straight_duct_T3_example

function allrun()
    println("#####################################################")
    println("# straight_duct_Q8_example ")
    straight_duct_Q8_example()
    println("#####################################################")
    println("# straight_duct_T3_example ")
    straight_duct_T3_example()
end # function allrun

end # module straight_duct_examples


using COSMO, Test, LinearAlgebra, Statistics, Random

tol = 1e-3
rng = Random.MersenneTwister(41)
mutable struct TestProblem
    P
    q
    constraints
end

function simpleQP()
    P = [4. 1;1 2]
    q = [1; 1.]
    A = [1. 1;1 0; 0 1]
    l = [1.;0;0]
    u = [1.;0.7;0.7]

    constraint1 = COSMO.Constraint(-A, u, COSMO.Nonnegatives)
    constraint2 = COSMO.Constraint(A, -l, COSMO.Nonnegatives)
    constraints = [constraint1; constraint2]
    return TestProblem(P, q, constraints)
end

@testset "Simple Tests" begin

    @testset "Simple QP" begin
        p = simpleQP()
        model = COSMO.Model()
        assemble!(model, p.P, p.q, p.constraints, settings = COSMO.Settings(verbose=true))

        res = COSMO.optimize!(model);

        @test res.status == :Solved
        @test isapprox(norm(res.x - [0.3; 0.7]), 0., atol=tol)
        @test isapprox(res.obj_val, 1.8800000298331538, atol=tol)

    end

    @testset "update_max_iter" begin
        p = simpleQP()
        settings = COSMO.Settings(max_iter=20)
        model = COSMO.Model()
        assemble!(model, p.P, p.q, p.constraints, settings = settings)

        res = COSMO.optimize!(model);
        @test res.status == :Max_iter_reached
    end

    @testset "update_check_termination" begin
         p = simpleQP()
        settings = COSMO.Settings(check_termination=100000)
        model = COSMO.Model()
        assemble!(model, p.P, p.q, p.constraints, settings = settings)

        res = COSMO.optimize!(model);

        @test res.status == :Max_iter_reached

    end

    @testset "update_rho" begin

        # TODO: write problem where rho is updated, results should be <exactly> the same

    end

    @testset "timelimit" begin
        p = simpleQP()
        settings = COSMO.Settings(time_limit=1, check_termination=100000000,max_iter=10000000)
        model = COSMO.Model()
        assemble!(model, p.P, p.q, p.constraints, settings = settings)

        res = COSMO.optimize!(model);
        @test res.status == :Time_limit_reached
    end

     @testset "warm_starting" begin
        p = simpleQP()
        settings = COSMO.Settings(check_termination = 1)
        model = COSMO.Model()
        assemble!(model, p.P, p.q, p.constraints, settings = settings)

        res1 = COSMO.optimize!(model);
        n = 2
        m = 6
        x0 = res1.x
        y0 = res1.y
        s0 = res1.s

        model = COSMO.Model()
        assemble!(model, p.P, p.q, p.constraints, settings = settings)
        COSMO.warm_start_primal!(model, x0)
        COSMO.warm_start_dual!(model, y0)
        COSMO.warm_start_slack!(model, s0)
        res2 = COSMO.optimize!(model);

        @test res1.status == :Solved && res2.status == :Solved && res2.iter < res1.iter

        @test_throws DimensionMismatch COSMO.warm_start_primal!(model, rand(4))
        @test_throws DimensionMismatch COSMO.warm_start_dual!(model, rand(2))
    end

end
nothing


struct MInfo{T}
    idx_q1::Vector{Int}
    idx_q2::Vector{Int}
    idx_u1::Vector{Int}
    v1::Vector{T}
end

function _step(sim::Simulator, info::MInfo, q1, q2, u1)
    info.v1 .= q2
    info.v1 .-= q1
    info.v1 ./= sim.h
    RoboDojo.step!(sim, q2, info.v1, u1, 1)
end

struct GradientBundle{T}
    ls::LeastSquares{T}
    q1η::Vector{T}
    q2η::Vector{T}
    u1η::Vector{T}
    dz::Matrix{T}
    ny::Int
    nz::Int
    info::MInfo{T}
end

function GradientBundle(model; N=100, ϵ=1.0e-4)
    nx = 2 * model.nq
    ny = model.nq
    nu = model.nu
    nz = nx + nu
    nθ = ny * (nz)

    @variables fz[1:ny] fη[1:ny] η[1:nz] θ[1:nθ]

    function cost(fz, fη, η, θ)
        M = reshape(θ, ny, nz)
        r = fη - fz - M * η
        return [transpose(r) * r]
    end

    c = cost(fz, fη, η, θ)
    cθ = Symbolics.gradient(c[1], θ)
    cθθ = Symbolics.hessian(c[1], θ)

    c_func_ls = eval(Symbolics.build_function(c, fz, fη, η, θ)[2])
    cθ_func_ls = eval(Symbolics.build_function(cθ, fz, fη, η, θ)[2])
    cθθ_func_ls = eval(Symbolics.build_function(cθθ, fz, fη, η, θ)[2])

    η = [zeros(nz) for i = 1:N]

    for i = 1:N
        w = ϵ * randn(1)[1]
        η[i][rand([j for j = 1:nz], 1)[1]] = w
    end

    # for i = 1:nz
    #     η[i][i] = ϵ
    #     η[i+nz][i] = -ϵ
    # end

    info = MInfo(
        collect(1:model.nq),
        collect(model.nq .+ (1:model.nq)),
        collect(2 * model.nq .+ (1:model.nu)),
        zeros(model.nq))

    fz = zeros(model.nq)
    fη = [zeros(model.nq) for i = 1:N]

    ls = LeastSquares(N,
        fz, fη, η,
        c_func_ls, cθ_func_ls, cθθ_func_ls,
        [0.0], zeros(nθ), zeros(nθ, nθ),
        [0.0], zeros(nθ), zeros(nθ, nθ),
        zeros(nθ), zeros(nθ),
        lu_solver(zeros(nθ, nθ)))

    GradientBundle(ls,
        zeros(nq),
        zeros(nq),
        zeros(nu),
        zeros(ny, nz),
        ny, nz,
        info)
end

function gradient!(sim::Simulator, gb::GradientBundle, q1, q2, u1; verbose=false)
    gb.ls.fz .= _step(sim, gb.info, q1, q2, u1)
    for i = 1:gb.ls.N
        ηq1 = @views gb.ls.η[i][gb.info.idx_q1]
        ηq2 = @views gb.ls.η[i][gb.info.idx_q2]
        ηu1 = @views gb.ls.η[i][gb.info.idx_u1]
        gb.q1η .= q1
        gb.q1η .+= ηq1
        gb.q2η .= q2
        gb.q2η .+= ηq2
        gb.u1η .= u1
        gb.u1η .+= ηu1
        gb.ls.fη[i] .= _step(sim, gb.info, gb.q1η, gb.q2η, gb.u1η)
    end
    update!(gb.ls, verbose=verbose)
    gb.dz .= Base.ReshapedArray(gb.ls.θ, (gb.ny, gb.nz), ())
    return gb.dz
end

# gradient!(im_dyn.eval_sim, gb, q1, q2, u1)
# @benchmark gradient!($im_dyn.eval_sim, $gb, $q1, $q2, $u1)

function fx_gb(dx, model::ImplicitDynamics, x, u, w)
	q1 = @views x[model.idx_q1]
	q2 = @views x[model.idx_q2]
	model.v1 .= q2
	model.v1 .-= q1
	model.v1 ./= model.eval_sim.h
	RoboDojo.step!(model.eval_sim, q2, model.v1, u, 1)
	nq = model.eval_sim.model.nq
	for i = 1:nq
		dx[model.idx_q1[i], model.idx_q2[i]] = 1.0
	end
    gradient!(model.eval_sim, model.info, q1, q2, u)
    ∂q3∂q1 = @views model.info.dz[model.idx_q1, model.idx_q1]
    ∂q3∂q2 = @views model.info.dz[model.idx_q1, model.idx_q2]
	dx[model.idx_q2, model.idx_q1] = ∂q3∂q1
	dx[model.idx_q2, model.idx_q2] = ∂q3∂q2
	return dx
end

# dx = zeros(2 * planarpush.nq, 2 * planarpush.nq)
# x = rand(2 * planarpush.nq)
# u = rand(planarpush.nu)
# w = zeros(0)
# fx_gb(dx, im_dyn, x, u, w)
# @benchmark fx_gb($dx, $im_dyn, $x, $u, $w)
# @code_warntype fx_gb(dx, im_dyn, x, u, w)

function fu_gb(du, model::ImplicitDynamics, x, u, w)
	q1 = @views x[model.idx_q1]
	q2 = @views x[model.idx_q2]
	model.v1 .= q2
	model.v1 .-= q1
	model.v1 ./= model.eval_sim.h
	RoboDojo.step!(model.eval_sim, q2, model.v1, u, 1)
    gradient!(model.eval_sim, model.info, q1, q2, u)
    ∂q3∂u1 = @views model.info.dz[model.idx_q1, model.info.info.idx_u1]
	du[model.idx_q2, :] = ∂q3∂u1
	return du
end

# du = zeros(2 * planarpush.nq, planarpush.nu)
# x = rand(2 * planarpush.nq)
# u = rand(planarpush.nu)
# w = zeros(0)
# fu_gb(du, im_dyn, x, u, w)
# @benchmark fu_gb($du, $im_dyn, $x, $u, $w)
# @code_warntype fu_gb(du, im_dyn, x, u, w)


"""
Generate a population of size `n` distributed spatially according to
a bivariate distribution, `d`
"""
function create_pop_db_bivariate(n::Int64, d::ContinuousMultivariateDistribution)
  coordinates = rand(d, n)'
  DataFrame(ind_id = 1:n, x = coordinates[:,1], y = coordinates[:,2])
end

"""
Generate a population of size `n` distributed spatially according to
a univariate distributions, `d1` and `d2`
"""
function create_pop_db_univariates(n::Int64,
                                   d1::ContinuousUnivariateDistribution,
                                   d2::ContinuousUnivariateDistribution)
  DataFrame(ind_id = 1:n, x = rand(d1, n), y = rand(d2, n))
end

"""
Construct a matrix containing the euclidean distance between every
pair of individuals contained in `pop_db`
"""
function create_dist_mtx(pop_db::DataFrame)
  distance_mat = fill(0.0, (pop_db[end, 1], pop_db[end, 1]))
  for i = 1:(size(distance_mat)[1]-1)
    for j = (i+1):(size(distance_mat)[1])
      distance_mat[i,j] = euclidean([pop_db[i, 2],pop_db[i, 3]], [pop_db[j, 2],pop_db[j, 3]])
    end
  end
  distance_mat' + distance_mat
end


module FractalFlow

using StaticArrays   # SVector
using IterTools      # Imap
using LightXML       # Parse XML, VTU
using LightGraphs    #
using Combinatorics  # Combinations
using DataStructures # Deque
using Polynomials
using WriteVTK
using PyCall        # need networkx on python side

import Base: length, write, size, split, start, next, done, scale!
import LightGraphs: nv, ne, Graph, degree
# Topology
export Curve, VertexFunction, EdgeFunction, is_consistent, entity_dim
export nv, ne, Graph, degree_f, num_loops, branch_map, num_branches, num_colors, color_f
# Geometry
export FractalDim, fractal_dim, estimate_fd, min_branch, max_branch
export min_segment, max_segment
# Generation, fractals
export koch_t_curve, koch_q1_curve, koch_q2_curve
# Generation, random
export gmsh_code
# IO
export write

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

"""
Curve is a collection of segments. Each segment is defined by 2 indices
corresponding to two points in the points array.
"""
struct Curve{D, T}
    # NOTE: The inner preconditioner assumes sensible data that is points
    # in the vector are unique and the segments are unique as well. We do
    # no check if this is true here
    points::Vector{SVector{D, T}}
    segments::Vector{Tuple{Int, Int}}
end
geometrical_dim{D, T}(::Curve{D, T}) = D

"""|AB|"""
struct Segment{D, T}
    A::SVector{D, T}
    B::SVector{D, T}
end

# Constructors
include("curve.jl")

abstract type CurveFunction end

geometrical_dim(f::CurveFunction) = geometrical_dim(f.curve)

"""Values in nodes"""
struct VertexFunction{R <: Real} <: CurveFunction
    curve::Curve
    data::Vector{R}

    function VertexFunction{R}(c::Curve, d::Vector{R}) where R <: Real
        @assert length(c.points) == length(d)
        new(c, d)
    end
end
VertexFunction(c::Curve, d::Vector) = VertexFunction{eltype(d)}(c, d)

topological_dim(::VertexFunction) = 0

"""Values in segments"""
struct EdgeFunction{R <: Real} <: CurveFunction
    curve::Curve
    data::Vector{R}

    function EdgeFunction{R}(c::Curve, d::Vector{R}) where R <: Real
        @assert length(c.segments) == length(d)
        new(c, d)
    end
end
EdgeFunction(c::Curve, d::Vector) = EdgeFunction{eltype(d)}(c, d)

topological_dim(::EdgeFunction) = 1

# IO - let curves be 'readable' from files
include("vtk_io.jl")
include("xml_io.jl")
include("msh_io.jl")

"""Load from VTU or XML file"""
function Curve(path::AbstractString)
    name, ext = splitext(path)

    ext == ".vtu" && return Curve(read_curve_vtu(path)...)

    ext == ".xml" && return Curve(read_curve_xml(path)...)

    ext == ".msh" && return Curve(read_curve_msh(path)...)
end

# Various properties of curves
include("topology.jl")  # Connectivity

include("geometry.jl")  # Distance, fractal properties

# Generation
include("generation.jl")  # Fractals and random curves
include("gmsh_code.jl")  #

end # module


using Agents, Random, DataFrames, LightGraphs
using Distributions: Poisson, DiscreteNonParametric
using DrWatson: @dict
using LinearAlgebra: diagind

mutable struct PoorSoul <: AbstractAgent
    id::Int
    pos::Int
    days_infected::Int  # number of days since is infected
    status::Symbol  # 1: S, 2: I, 3:R
end

function model_initiation(;
    Ns = [500, 500, 500],
    migration_rate = 0.2,
    death_rate = 0.1,
    β_det = 0.05,
    β_und = 0.1,
    infection_period = 10,
    reinfection_probability = 0.05,
    detection_time = 3,
    Is = ones(Int, length(Ns)),
    seed = nothing,
)

    if !isnothing(seed)
        Random.seed!(seed)
    end

    C = length(Ns)
    β_det = [β_det for i in 1:C]
    β_und = [β_und for i in 1:C]

    migration_rates = reshape([migration_rate for i in 1:(C * C)], C, C)
    # normalize migration_rates
    migration_rates_sum = sum(migration_rates, dims = 2)
    for c in 1:C
        migration_rates[c, :] ./= migration_rates_sum[c]
    end

    properties = @dict(
        Ns,
        β_und,
        β_det,
        migration_rates,
        infection_period,
        reinfection_probability,
        detection_time,
        death_rate,
        Is,
        C,
    )
    space = GraphSpace(complete_digraph(C))
    model = ABM(PoorSoul, space; properties)

    # Add initial individuals
    for city in 1:C, n in 1:Ns[city]
        add_agent!(city, model, 0, :S) # Susceptible
    end
    # add infected individuals
    for city in 1:C
        inds = ids_in_position(city, model)
        for n in 1:Is[city]
            agent = model[inds[n]]
            agent.status = :I # Infected
            agent.days_infected = 1
        end
    end
    return model
end

function agent_step!(agent, model)
    migrate!(agent, model)
    transmit!(agent, model)
    update_status!(agent, model)
    recover_or_die!(agent, model)
end

function migrate!(agent, model)
    pid = agent.pos
    d = DiscreteNonParametric(1:(model.C), model.migration_rates[pid, :])
    m = rand(model.rng, d)
    if m ≠ pid
        move_agent!(agent, m, model)
    end
end

function transmit!(agent, model)
    agent.status == :S && return
    rate = if agent.days_infected < model.detection_time
        model.β_und[agent.pos]
    else
        model.β_det[agent.pos]
    end

    d = Poisson(rate)
    n = rand(model.rng, d)
    n == 0 && return

    for contactID in ids_in_position(agent, model)
        contact = model[contactID]
        if contact.status == :S ||
           (contact.status == :R && rand(model.rng) ≤ model.reinfection_probability)
            contact.status = :I
            n -= 1
            n == 0 && return
        end
    end
end

update_status!(agent, model) = agent.status == :I && (agent.days_infected += 1)

function recover_or_die!(agent, model)
    if agent.days_infected ≥ model.infection_period
        if rand(model.rng) ≤ model.death_rate
            kill_agent!(agent, model)
        else
            agent.status = :R
            agent.days_infected = 0
        end
    end
end


"""
Learning repo for testing various CI and benchmarking tools.

Mainly benchmarking my naive ForwardDiff implementation and Dormand-Prince solver on models I find interesting.

Working through `Mathematical Physiology` by James Keener and `Biomolecular Feedback Systems` by Murray.
"""
module DemoActions

using StaticArrays

include("mysolvers.jl")
include("models.jl")
include("fwd.jl")

# solver interfaces
export MyProbInplace, mysolve, mysolve_inplace

# solver algorithms
export eulers_inplace_f!,
    midpoint!,
    midpoint_inplace_f!,
    rk4!,
    rk4_inplace_f!,
    rk4_inplace_f_prealloc_k!,
    dp5,
    dp5!,
    dp5_inplace!

# mini model-zoo
export lotka!, transcriptional_component!, bistable_gene_circuit!, negative_autoregulation!, phosphorylation_cycle!

end


# `permutedims` seems to work better with AD (cf. KernelFunctions.jl)
aT_b(a::AbstractVector{<:Real}, b::AbstractMatrix{<:Real}) = permutedims(a) * b
# `permutedims` can't be used here since scalar output is desired
aT_b(a::AbstractVector{<:Real}, b::AbstractVector{<:Real}) = dot(a, b)

# flatten arrays with fallback for scalars
_vec(x::AbstractArray{<:Real}) = vec(x)
_vec(x::Real) = x


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

# %%
a = 1.3; tmax = 25.0; using Distributions, DifferentialEquations, Plots
kuramoto!(dθ, θ, param, t) = (N = length(θ); (K, ω) = param;
    for i in 1:N dθ[i] = ω[i] + K*mean(sin(θ[j] - θ[i]) for j in 1:N) end)
(m, n) = (32, 16); (d, v) = (Normal(0, 2), 1.0); K_c = 2/π/pdf(d, 0)
θ₀ = 2π*rand(m, n); tspan = (0.0, tmax); K = a*K_c; ω = rand(d, m, n) .+ v
sol = solve(ODEProblem(kuramoto!, θ₀, tspan, (K, ω)))
anim = @animate for t in [fill(0, 20); 0:0.1:tmax; fill(tmax, 20)]
    plot(size=(400, 220), title="Kuramoto model K = $(a)K_c: t = $t", titlefontsize=10)
    heatmap!(sin.(sol(t))'; c=:bwr, colorbar=false, frame=false, ticks=false)
end; gif(anim, "kuramoto$(a).gif")

# %%
quote
a = 1.3; tmax = 25.0; using Distributions, DifferentialEquations, Plots
kuramoto!(dθ, θ, param, t) = (N = length(θ); (K, ω) = param;
    for i in 1:N dθ[i] = ω[i] + K*mean(sin(θ[j] - θ[i]) for j in 1:N) end)
(m, n) = (32, 16); (d, v) = (Normal(0, 2), 1.0); K_c = 2/π/pdf(d, 0)
θ₀ = 2π*rand(m, n); tspan = (0.0, tmax); K = a*K_c; ω = rand(d, m, n) .+ v
sol = solve(ODEProblem(kuramoto!, θ₀, tspan, (K, ω)))
anim = @animate for t in [fill(0, 20); 0:0.1:tmax; fill(tmax, 20)]
    plot(size=(400, 220), title="Kuramoto model K = $(a)K_c: t = $t", titlefontsize=10)
    heatmap!(sin.(sol(t))'; c=:bwr, colorbar=false, frame=false, ticks=false)
end; gif(anim, "kuramoto$(a).gif")
end |> Base.remove_linenums! |> x -> display("text/markdown", "```julia\n$x\n```")

# %%
begin
    a = 1.3
    tmax = 25.0
    using Distributions, DifferentialEquations, Plots
    kuramoto!(dθ, θ, param, t) = begin
            N = length(θ)
            (K, ω) = param
            for i = 1:N
                dθ[i] = ω[i] + K * mean((sin(θ[j] - θ[i]) for j = 1:N))
            end
        end
    (m, n) = (32, 16)
    (d, v) = (Normal(0, 2), 1.0)
    K_c = (2 / π) / pdf(d, 0)
    θ₀ = (2π) * rand(m, n)
    tspan = (0.0, tmax)
    K = a * K_c
    ω = rand(d, m, n) .+ v
    sol = solve(ODEProblem(kuramoto!, θ₀, tspan, (K, ω)))
    anim = @animate(for t = [fill(0, 20); 0:0.1:tmax; fill(tmax, 20)]
                plot(size = (400, 220), title = "Kuramoto model K = $(a)K_c: t = $(t)", titlefontsize = 10)
                heatmap!((sin.(sol(t)))'; c = :bwr, colorbar = false, frame = false, ticks = false)
            end)
    gif(anim, "kuramoto$(a).gif")
end

# %%