## Common implementation of each block:
## returns a tuple of (output, internal_derivative)

# Low pass modified with denominator K_den instead of 1.0
"""
Low Pass Filter Modified
     ┌─────────────┐
     │      K      │
u -> │ ────────────│ -> y
     │ K_den + sT  │
     └─────────────┘
"""
function low_pass_modified_mass_matrix(
    u::Z,
    y::V,
    K::Float64,
    K_den::W,
    ::Float64,
) where {V <: ACCEPTED_REAL_TYPES, W <: ACCEPTED_REAL_TYPES, Z <: ACCEPTED_REAL_TYPES}
    return y, K * u - K_den * y
end

function low_pass_modified(
    u::Z,
    y::V,
    K::Float64,
    K_den::W,
    T::Float64,
) where {V <: ACCEPTED_REAL_TYPES, W <: ACCEPTED_REAL_TYPES, Z <: ACCEPTED_REAL_TYPES}
    return y, (1.0 / T) * low_pass_modified_mass_matrix(u, y, K, K_den, T)[2]
end

"""
Low Pass Filter
     ┌────────┐
     │    K   │
u -> │ ────── │ -> y
     │ 1 + sT │
     └────────┘
"""

# Use this one if T = 0 is allowed, and let the mass matrix take care of it.
function low_pass_mass_matrix(
    u::Z,
    y::V,
    K::Float64,
    T::Float64,
) where {V <: ACCEPTED_REAL_TYPES, Z <: ACCEPTED_REAL_TYPES}
    return low_pass_modified_mass_matrix(u, y, K, 1.0, T)
end

function low_pass(
    u::Z,
    y::V,
    K::Float64,
    T::Float64,
) where {V <: ACCEPTED_REAL_TYPES, Z <: ACCEPTED_REAL_TYPES}
    return low_pass_modified(u, y, K, 1.0, T)
end

"""
Low Pass Filter with non-windup
             y_max
           /¯¯¯¯¯¯
     ┌────────┐
     │    K   │
u -> │ ────── │ -> y
     │ 1 + sT │
     └────────┘
   ______/
   y_min
"""

function low_pass_nonwindup_mass_matrix(
    u::Z,
    y::V,
    K::Float64,
    ::Float64,
    y_min::Float64,
    y_max::Float64,
) where {V <: ACCEPTED_REAL_TYPES, Z <: ACCEPTED_REAL_TYPES}
    dydt = K * u - y
    y_sat = clamp(y, y_min, y_max)
    # Non Windup logic from IEEE Std 421.5
    binary_logic = ((y >= y_max) && (dydt > 0)) || ((y <= y_min) && (dydt < 0)) ? 0.0 : 1.0
    return y_sat, binary_logic * dydt
end

# Does not accept T = 0
function low_pass_nonwindup(
    u::Z,
    y::V,
    K::Float64,
    T::Float64,
    y_min::Float64,
    y_max::Float64,
) where {V <: ACCEPTED_REAL_TYPES, Z <: ACCEPTED_REAL_TYPES}
    y_sat, dydt_scaled = low_pass_nonwindup_mass_matrix(u, y, K, T, y_min, y_max)
    return y_sat, (1.0 / T) * dydt_scaled
end

# Does not accept T = 0
function low_pass_nonwindup_ramp_limits(
    u::Z,
    y::V,
    K::Float64,
    T::Float64,
    y_min::Float64,
    y_max::Float64,
    dy_min::Float64,
    dy_max::Float64,
) where {V <: ACCEPTED_REAL_TYPES, Z <: ACCEPTED_REAL_TYPES}
    y, dydt = low_pass(u, y, K, T)
    y_sat = clamp(y, y_min, y_max)
    dydt_sat = clamp(dydt, dy_min, dy_max)
    # Non Windup logic from IEEE Std 421.5
    binary_logic =
        ((y >= y_max) && (dydt_sat > 0)) || ((y <= y_min) && (dydt_sat < 0)) ? 0.0 : 1.0
    return y_sat, binary_logic * dydt_sat
end

"""
High Pass Filter
     ┌────────┐
     │   sK   │
u -> │ ────── │ -> y
     │ 1 + sT │
     └────────┘

Internal State: x
"""

function high_pass_mass_matrix(
    u::Z,
    x::Z,
    K::Float64,
    T::Float64,
) where {Z <: ACCEPTED_REAL_TYPES}
    K_T = T < eps() ? 0.0 : (K / T)
    return x + K_T * u, -(K_T * u + x)
end

function high_pass(u::Z, x::Z, K::Float64, T::Float64) where {Z <: ACCEPTED_REAL_TYPES}
    y, dxdt_scaled = high_pass_mass_matrix(u, x, K, T)
    return y, (1.0 / T) * dxdt_scaled
end

"""
Lead-Lag Block:
     ┌───────────┐
     │   1 + sT1 │
u -> │ K ─────── │ -> y
     │   1 + sT2 │
     └───────────┘

Internal State: x
"""

function lead_lag_mass_matrix(
    u::Z,
    x::Z,
    K::Float64,
    T1::Float64,
    T2::Float64,
) where {Z <: ACCEPTED_REAL_TYPES}
    T1_T2 = T2 < eps() ? 0.0 : (T1 / T2)
    return x + (K * T1_T2) * u, K * (1 - T1_T2) * u - x
end

function lead_lag(
    u::Z,
    x::Z,
    K::Float64,
    T1::Float64,
    T2::Float64,
) where {Z <: ACCEPTED_REAL_TYPES}
    y, dxdt_scaled = lead_lag_mass_matrix(u, x, K, T1, T2)
    return y, (1.0 / T2) * dxdt_scaled
end

"""
2nd Order Low Pass Filter
     ┌──────────────────┐
     │         K        │
u -> │ ──────────────── │ -> y
     │ 1 + sT1 + s^2 T2 │
     └──────────────────┘

Internal state: x
"""

function low_pass_2nd(
    u::Z,
    x::Z,
    y::Z,
    K::Float64,
    T1::Float64,
    T2::Float64,
) where {Z <: ACCEPTED_REAL_TYPES}
    dxdt = (1.0 / T2) * (K * u - T1 * x - y)
    dydt = x
    return y, dxdt, dydt
end

function low_pass_2nd_mass_matrix(
    u::Z,
    x::Z,
    y::Z,
    K::Float64,
    T1::Float64,
    T2::Float64,
) where {Z <: ACCEPTED_REAL_TYPES}
    dxdt = K * u - T1 * x - y
    dydt = x
    return y, dxdt, dydt
end

"""
2nd Order Lead-Lag Block
     ┌──────────────────┐
     │ 1 + sT3 + s^2 T4 │
u -> │ ──────────────── │ -> y
     │ 1 + sT1 + s^2 T2 │
     └──────────────────┘

Internal States: x1, x2
"""

function lead_lag_2nd_mass_matrix(
    u::Z,
    x1::Z,
    x2::Z,
    T1::Float64,
    T2::Float64,
    T3::Float64,
    T4::Float64,
) where {Z <: ACCEPTED_REAL_TYPES}
    dx1dt = u - T1 * x1 - x2
    dx2dt = x1
    T4_T2 = T2 < eps() ? 0.0 : (T4 / T2)
    y = T4_T2 * u + (T3 - T1 * T4_T2) * x1 + (1.0 - T4_T2) * x2
    return y, dx1dt, dx2dt
end

function lead_lag_2nd(
    u::Z,
    x1::Z,
    x2::Z,
    T1::Float64,
    T2::Float64,
    T3::Float64,
    T4::Float64,
) where {Z <: ACCEPTED_REAL_TYPES}
    y, dx1dt_scaled, dx2dt = lead_lag_2nd_mass_matrix(u, x1, x2, T1, T2, T3, T4)
    return y, (1.0 / T2) * dx1dt_scaled, dx2dt
end

"""
Proportional-Integral Block
             y_max
            /¯¯¯¯¯¯
     ┌──────────┐
     │      ki  │
u -> │kp + ───  │ -> y
     │      s   │
     └──────────┘
    ______/
     y_min

Internal State: x
"""

function pi_block(u::Z, x::Z, kp::Float64, ki::Float64) where {Z <: ACCEPTED_REAL_TYPES}
    return kp * u + ki * x, u
end

function pi_block_nonwindup(
    u::Z,
    x::Z,
    kp::Float64,
    ki::Float64,
    y_min::Float64,
    y_max::Float64,
) where {Z <: ACCEPTED_REAL_TYPES}
    y, _ = pi_block(u, x, kp, ki)
    y_sat = clamp(y, y_min, y_max)
    binary_logic = y_min < y < y_max ? 1.0 : 0.0
    return y_sat, binary_logic * u
end

"""
Integrator with windup limits
                                 y_max
                                _ _ _
        ┌────────┐             /
        │    K   │   y        /
 u - -->│ ────── │ - - - - - /- - - - - --> ysat
        │   sT   │          /
        └────────┘   _ _ _ /
                     y_min
 """

function integrator_windup_mass_matrix(
    u::Z,
    y::Z,
    K::Float64,
    ::Float64,
    y_min::Float64,
    y_max::Float64,
) where {Z <: ACCEPTED_REAL_TYPES}
    dydt_scaled = K * u
    y_sat = clamp(y, y_min, y_max)
    return y_sat, dydt_scaled
end

# Does not accept T = 0
function integrator_windup(
    u::Z,
    y::Z,
    K::Float64,
    T::Float64,
    y_min::Float64,
    y_max::Float64,
) where {Z <: ACCEPTED_REAL_TYPES}
    y_sat = clamp(y, y_min, y_max)
    return y_sat, (1.0 / T) * integrator_windup_mass_matrix(u, y, K, T, y_min, y_max)[2]
end

"""
Integrator with non-windup limits
             y_max
           /¯¯¯¯¯¯
     ┌────────┐
     │    K   │
u -> │ ────── │ -> y
     │   sT   │
     └────────┘
   ______/
   y_min
"""

function integrator_nonwindup_mass_matrix(
    u::Z,
    y::Z,
    K::Float64,
    ::Float64,
    y_min::Float64,
    y_max::Float64,
) where {Z <: ACCEPTED_REAL_TYPES}
    dydt_scaled = K * u
    y_sat = clamp(y, y_min, y_max)
    upper_lim = (y >= y_max) && (dydt_scaled > 0)
    lower_lim = (y <= y_min) && (dydt_scaled < 0)
    binary_logic = upper_lim || lower_lim ? 0.0 : 1.0
    return y_sat, binary_logic * dydt_scaled
end

# Does not accept T = 0
function integrator_nonwindup(
    u::Z,
    y::Z,
    K::Float64,
    T::Float64,
    y_min::Float64,
    y_max::Float64,
) where {Z <: ACCEPTED_REAL_TYPES}
    y_sat, dydt_scaled = integrator_nonwindup_mass_matrix(u, y, K, T, y_min, y_max)
    return y_sat, (1.0 / T) * dydt_scaled
end


using Gadfly
using Compose
using Distributions
import Cairo, Fontconfig
using DataFrames

# %%
g = Geometric(1/2)

# %%
x_min = [0 1/2 3/4 7/8 15/16]
x_max = [1/2 3/4 7/8 15/16 31/32]
y = [2 4 8 16 32]

label = ["H" "TH" "TTH" "..."]
text_x = [1/4 5/8 13/16 29/32]

tick_labels=Dict(zip(x_min, ["0" "1/2" "1/4" "1/8" "1/16"]))

theme = Theme(major_label_font="PragmataPro Mono",
              minor_label_font="PragmataPro Mono",
              key_title_font="PragmataPro Mono Bold",
              highlight_width=0.005w,
              bar_highlight=color("black"))
Gadfly.push_theme(theme)

p = plot(
    x_min=x_min,
    x_max=x_max,
    y=y,
    Geom.bar,
    Coord.cartesian(ymax=20),
    Guide.xticks(ticks=x_min),
    Guide.yticks(ticks=y),
    Scale.x_continuous(labels = x -> tick_labels[x]),
    Guide.xlabel("PROBABILITY"),
    Guide.ylabel("PAYOUT"),
    Guide.title("<b>Distribution of Payouts</b>"),
    Guide.annotation(compose(
        context(),
        text(text_x, ones(5), label, repeat([hcenter], 4), repeat([vcenter], 4)),
        # font("PragmataPro Mono"),
        fontsize(8pt),
        stroke("black"))))
# %%
draw(PDF("../../figures/distribution.pdf"), p)


@testset "ucme.jl" begin
    @testset "UCME: Binary examples" begin
        # categorical distributions
        ucme = UCME(
            SqExponentialKernel() ⊗ WhiteKernel(),
            [[1.0, 0], [0.5, 0.5], [0.0, 1]],
            [1, 1, 2],
        )
        for predictions in ([[1, 0], [0, 1]], ColVecs([1 0; 0 1]), RowVecs([1 0; 0 1]))
            @test iszero(@inferred(ucme(predictions, [1, 2])))
            @test @inferred(ucme(predictions, [1, 1])) ≈ (exp(-2) + exp(-0.5) + 1) / 12
            @test @inferred(ucme(predictions, [2, 1])) ≈ (1 - exp(-1))^2 / 6
            @test @inferred(ucme(predictions, [2, 2])) ≈ (exp(-2) + exp(-0.5) + 1) / 12
        end

        # probabilities
        ucme = UCME(
            (SqExponentialKernel() ∘ ScaleTransform(sqrt(2))) ⊗ WhiteKernel(),
            [1.0, 0.5, 0.0],
            [true, true, false],
        )
        @test iszero(@inferred(ucme([1, 0], [true, false])))
        @test @inferred(ucme([1, 0], [true, true])) ≈ (exp(-2) + exp(-0.5) + 1) / 12
        @test @inferred(ucme([1, 0], [false, true])) ≈ (1 - exp(-1))^2 / 6
        @test @inferred(ucme([1, 0], [false, false])) ≈ (exp(-2) + exp(-0.5) + 1) / 12
    end

    @testset "UCME: Basic properties" begin
        estimates = Vector{Float64}(undef, 1_000)

        for ntest in (1, 5, 10)
            # categorical distributions
            for nclasses in (2, 10, 100)
                dist = Dirichlet(nclasses, 1.0)

                testpredictions = [rand(dist) for _ in 1:ntest]
                testtargets = rand(1:nclasses, ntest)
                ucme = UCME(
                    (ExponentialKernel() ∘ ScaleTransform(0.1)) ⊗ WhiteKernel(),
                    testpredictions,
                    testtargets,
                )

                predictions = [Vector{Float64}(undef, nclasses) for _ in 1:20]
                targets = Vector{Int}(undef, 20)

                for i in 1:length(estimates)
                    rand!.(Ref(dist), predictions)
                    targets .= rand.(Categorical.(predictions))

                    estimates[i] = ucme(predictions, targets)
                end

                @test all(x > zero(x) for x in estimates)
            end

            # probabilities
            testpredictions = rand(ntest)
            testtargets = rand(Bool, ntest)
            ucme = UCME(
                (ExponentialKernel() ∘ ScaleTransform(0.1)) ⊗ WhiteKernel(),
                testpredictions,
                testtargets,
            )

            predictions = Vector{Float64}(undef, 20)
            targets = Vector{Bool}(undef, 20)

            for i in 1:length(estimates)
                rand!(predictions)
                map!(targets, predictions) do p
                    return rand() < p
                end

                estimates[i] = ucme(predictions, targets)
            end

            @test all(x > zero(x) for x in estimates)
        end
    end

    # alternative implementation of white kernel
    struct WhiteKernel2 <: Kernel end
    (::WhiteKernel2)(x, y) = x == y

    # alternative implementation TensorProductKernel
    struct TensorProduct2{K1<:Kernel,K2<:Kernel} <: Kernel
        kernel1::K1
        kernel2::K2
    end
    function (kernel::TensorProduct2)((x1, x2), (y1, y2))
        return kernel.kernel1(x1, y1) * kernel.kernel2(x2, y2)
    end

    @testset "binary classification" begin
        # probabilities and boolean targets
        p, testp = rand(2)
        y, testy = rand(Bool, 2)
        scale = rand()
        kernel = SqExponentialKernel() ∘ ScaleTransform(scale)
        val = unsafe_ucme_eval(kernel ⊗ WhiteKernel(), p, y, testp, testy)
        @test unsafe_ucme_eval(kernel ⊗ WhiteKernel2(), p, y, testp, testy) ≈ val
        @test unsafe_ucme_eval(TensorProduct2(kernel, WhiteKernel()), p, y, testp, testy) ≈
              val
        @test unsafe_ucme_eval(TensorProduct2(kernel, WhiteKernel2()), p, y, testp, testy) ≈
              val

        # corresponding values and kernel for full categorical distribution
        pfull = [p, 1 - p]
        yint = 2 - y
        testpfull = [testp, 1 - testp]
        testyint = 2 - testy
        kernelfull = SqExponentialKernel() ∘ ScaleTransform(scale / sqrt(2))

        @test unsafe_ucme_eval(
            kernelfull ⊗ WhiteKernel(), pfull, yint, testpfull, testyint
        ) ≈ val
        @test unsafe_ucme_eval(
            kernelfull ⊗ WhiteKernel2(), pfull, yint, testpfull, testyint
        ) ≈ val
        @test unsafe_ucme_eval(
            TensorProduct2(kernelfull, WhiteKernel()), pfull, yint, testpfull, testyint
        ) ≈ val
        @test unsafe_ucme_eval(
            TensorProduct2(kernelfull, WhiteKernel2()), pfull, yint, testpfull, testyint
        ) ≈ val
    end

    @testset "multi-class classification" begin
        n = 10
        p = rand(n)
        p ./= sum(p)
        y = rand(1:n)
        testp = rand(n)
        testp ./= sum(testp)
        testy = rand(1:n)

        kernel = SqExponentialKernel() ∘ ScaleTransform(rand())
        val = unsafe_ucme_eval(kernel ⊗ WhiteKernel(), p, y, testp, testy)

        @test unsafe_ucme_eval(kernel ⊗ WhiteKernel2(), p, y, testp, testy) ≈ val
        @test unsafe_ucme_eval(TensorProduct2(kernel, WhiteKernel()), p, y, testp, testy) ≈
              val
        @test unsafe_ucme_eval(TensorProduct2(kernel, WhiteKernel2()), p, y, testp, testy) ≈
              val
    end
end


using JuMP

# Inspired from Joey Huchette's test in ConvexHull.jl
function jumpsimplextest(lib::PolyhedraLibrary, n)
    m = Model()
    @variable(m, x[1:n] >= 0)
    @constraint(m, sum(x[i] for i=1:n) == 1)

    poly = polyhedron(m, lib)

    V = zeros(Int, n, n)
    for k in 1:n
        V[k, k] = 1
    end
    generator_fulltest(poly, V)
end

# Inspired from Joey Huchette's test in ConvexHull.jl
function simplexorigtest(lib::PolyhedraLibrary, n)
    m = Model()
    @variable(m, x[1:n] >= 0)
    @constraint(m, sum(x[i] for i=1:n) ≤ 1)

    poly = polyhedron(m, lib)

    V = zeros(Int, n+1, n)
    for k in 1:n
        V[k, k] = 1
    end
    generator_fulltest(poly, V)
end


using JULES, Oceananigans
using Plots
using Printf

Nx = Ny = 1
Nz = 32
L = 10e3

grid = RegularCartesianGrid(size=(Nx, Ny, Nz), halo=(2, 2, 2),
                            x=(0, L), y=(0, L), z=(0, L))
#####
##### Dry adiabatic atmosphere
#####

gas = IdealGas()
Rᵈ, cₚ, cᵥ = gas.Rᵈ, gas.cₚ, gas.cᵥ

g  = 9.80665
pₛ = 100000
θₛ = 300
πₛ = 1

θ(x, y, z) = θₛ
π(x, y, z) = πₛ - g*z / (cₚ*θₛ)
p(x, y, z) = pₛ * π(x, y, z)^(cₚ/Rᵈ)
ρ(x, y, z) = pₛ / (Rᵈ * θ(x, y, z)) * π(x, y, z)^(cᵥ/Rᵈ)
Θ(x, y, z) = ρ(x, y, z) * θ(x, y, z)

#####
##### Create model
#####

model = CompressibleModel(grid=grid, buoyancy=gas, reference_pressure=pₛ,
                          thermodynamic_variable=ModifiedPotentialTemperature(),
                          tracers=(:Θᵐ,))

#####
##### Set initial conditions
#####

set!(model.density, ρ)
set!(model.tracers.Θᵐ, θ)

# Initial profiles
Θ₀_prof = model.tracers.Θᵐ[1, 1, 1:Nz]
ρ₀_prof = model.density[1, 1, 1:Nz]

# Arrays we will use to store a time series of ρw(z = L/2).
times = [model.clock.time]
ρw_ts = [model.momenta.ρw[1, 1, Int(Nz/2)]]

#####
##### Time step and keep plotting vertical profiles of ρθ′, ρw, and ρ′.
#####

# @animate for i=1:100
while model.clock.time < 500
    time_step!(model; Δt=0.5, Nt=10)

    @show model.clock.time
    Θ_prof = model.tracers.Θᵐ[1, 1, 1:Nz] .- Θ₀_prof
    W_prof = model.momenta.ρw[1, 1, 1:Nz+1]
    ρ_prof = model.density[1, 1, 1:Nz] .- ρ₀_prof

    Θ_plot = plot(Θ_prof, grid.zC, xlim=(-2.5, 2.5), xlabel="Theta_prime", ylabel="z (m)", label="")
    W_plot = plot(W_prof, grid.zF, xlim=(-1.2, 1.2), xlabel="rho*w", label="")
    ρ_plot = plot(ρ_prof, grid.zC, xlim=(-0.01, 0.01), xlabel="rho", label="")

    t_str = @sprintf("t = %d s", model.clock.time)
    display(plot(Θ_plot, W_plot, ρ_plot, title=["" "$t_str" ""], layout=(1, 3), show=true))

    push!(times, model.clock.time)
    push!(ρw_ts, model.momenta.ρw[1, 1, Int(Nz/2)])
end

#####
##### Plot the initial profile and the hydrostatically balanced profile.
#####

ρ∞_prof = model.density[1, 1, 1:Nz]

θ₀_prof = Θ₀_prof ./ ρ₀_prof
θ∞_prof = model.tracers.Θᵐ[1, 1, 1:Nz] ./ ρ∞_prof

θ_plot = plot(θ₀_prof, grid.zC, xlabel="theta (K)", ylabel="z (m)", label="initial", legend=:topleft)
plot!(θ_plot, θ∞_prof, grid.zC, label="balanced")

ρ_plot = plot(ρ₀_prof, grid.zC, xlabel="rho (kg/m³)", ylabel="z (m)", label="initial")
plot!(ρ_plot, ρ∞_prof, grid.zC, label="balanced")

display(plot(θ_plot, ρ_plot, show=true))

#####
##### Plot timeseries of maximum ρw
#####

plot(times, ρw_ts, linewidth=2, xlabel="time (s)", ylabel="rho*w(z=$(Int(L/2)) m)", label="", show=true)


using Symbolics
using StaticArrays

"""
The potential and number of dimentions of a bead
"""
struct BeadDefinition{Ndims}
    chain_pe
    chain_force
    perbead_param_keys
    global_param_keys
    VNdims::Val{Ndims}
    chain_pe_expression
end

function BeadDefinition(chain_pe, perbead_param_keys::NTuple{Nbparams,Symbol}, global_param_keys::NTuple{Ngparams,Symbol} , VNdims::Val{Ndims}) where {Ndims, Nbparams, Ngparams}
    #symbolically diff potential energy to get force
    @variables p[1:Ndims]
    @variables b[1:Ndims]
    @variables n[1:Ndims]
    @variables bpars[1:Nbparams]
    @variables gpars[1:Ngparams]
    bsymparams= (; zip(perbead_param_keys,(bpars...,))...)
    gsymparams= (; zip(global_param_keys,(gpars...,))...)
    exprpe= chain_pe(p,b,n,bsymparams,gsymparams)
    g= Symbolics.gradient(-exprpe, [p;b;n]; simplify=true)
    @debug g
    gexpr= build_function(g, [p;b;n], bpars, gpars ,expression=Val{true})
    gfun= eval(gexpr[1])
    function chain_force(p,b,n, bpars, gpars)
        f= gfun([p;b;n], bpars, gpars)
        pf = ntuple((i -> @inbounds(f[i])), VNdims)
        bf = ntuple((i -> @inbounds(f[i+Ndims])), VNdims)
        nf = ntuple((i -> @inbounds(f[i+2Ndims])), VNdims)
        (SVector(pf),SVector(bf),SVector(nf))
    end
    BeadDefinition(chain_pe, chain_force, perbead_param_keys, global_param_keys, VNdims, exprpe)
end

struct Chain
    beaddef::BeadDefinition
    perbead_params
    global_params
    init_pos
    init_vel
    function Chain(beaddef::BeadDefinition, perbead_params, global_params, init_pos, init_vel)
        #error checking
        size(init_pos) == size(init_vel) || throw(DimensionMismatch("init_pos $(size(init_pos)) and init_vel $(size(init_vel)) don't match"))
        length(size(init_pos)) == 2 || throw(DimensionMismatch("init_pos $(length(size(init_pos))) must be 2"))
        Val(size(init_pos)[2]) == beaddef.VNdims || throw(DimensionMismatch("init_pos $(size(init_pos)) must match bead dimension $(beaddef.VNdims)"))
        Nbeads= size(init_pos)[1]
        Nbeads == length(perbead_params) || throw(DimensionMismatch("number of coords $(Nbeads) must match number of params $(length(perbead_params))"))
        #make sure forces are finite
        for beadi in 1:Nbeads
            p= init_pos[mod(beadi-2,Nbeads)+1,:]
            b= init_pos[beadi,:]
            n= init_pos[mod(beadi,Nbeads)+1,:]
            pf, bf, nf = beaddef.chain_force(p,b,n,perbead_params[beadi], global_params)
            e= beaddef.chain_pe(p,b,n,perbead_params[beadi], global_params)
            all(isfinite.([pf;bf;nf;e])) || @warn "some forces or energies are not finite on Bead number $(beadi)"
            all(abs.([pf;bf;nf;e]) .< 2.0^30) || @warn "some forces or energies are over 2^30 on Bead number $(beadi)"
        end
        new(beaddef::BeadDefinition, copy(perbead_params), global_params, copy(init_pos), copy(init_vel))
    end
end

"""
zero velocities by default
"""
function Chain(beaddef::BeadDefinition, perbead_params, global_params, init_pos)
    Chain(beaddef::BeadDefinition, perbead_params, global_params, init_pos, zero(init_pos))
end

"""
A force that acts on a specific list of beads
"""
struct SpecificForce{Ndims,Nways}
    pe
    force
    params
    interactions::AbstractArray{NTuple{Nways,NTuple{2,T}},1} where T<:Integer
    VNdims::Val{Ndims}
    pe_expression
end

"""
Create a specific force
pe

"""
function SpecificForce(pe, params,
    interactions::AbstractArray{NTuple{Nways,NTuple{2,T}},1} where T<:Integer,
    VNdims::Val{Ndims}) where {Ndims, Nways}
    #symbolically diff potential energy to get force
    a = Num(SymbolicUtils.Sym{Real}(:a))
    @variables r[1:Nways,1:Ndims]
    @variables pars[1:length(params)]
    @variables time
    symparams= (; zip(keys(params),(pars...,))...)
    exprpe= pe(ntuple(i->r[i,:],Nways), time ,symparams) + 0time#This 0time is to force exprpe into a Num type
    rssymb= vcat(ntuple(i->r[i,:],Nways)...)
    g= Symbolics.gradient(-exprpe, rssymb; simplify=true)
    @debug g
    gexpr= build_function(g, rssymb, time, pars, expression=Val{true})
    gfun= eval(gexpr[1])
    function force(rs, time, pars)
        f= gfun(vcat(rs...), time, pars)
        ntuple(j->SVector(ntuple(i->f[i+Ndims*(j-1)],Ndims)),Nways)
    end
    SpecificForce(pe,force,params,copy(interactions),VNdims,exprpe)
end

"""
A collection of Chains and SpecificForces
"""
struct ChainSystem
    starttime
    sp_pes
    sp_forces
    sp_interactions
    sp_params
    beaddef::BeadDefinition
    perbead_params
    global_params
    init_pos
    init_vel
    """
    chainbounds is a tuple of at least 1 Int,
        chainbounds[1] is 1
        chainbounds[end] is the total number of beads + 1
        the other element in chainbounds are the bead ids of the start of a new chain
    """
    chainbounds::Tuple{T,Vararg{T}} where T <: Integer
end

"""
Create an empty System
"""
function ChainSystem(starttime, beaddef::BeadDefinition, global_params)
    ChainSystem(starttime,(),(),(),(),
                    beaddef,
                    nothing,global_params,
                    nothing,nothing,(1,))
end

"""
return a new ChainSystem with added chain
"""
function append(s::ChainSystem, c::Chain)
    #error checking
    s.beaddef == c.beaddef || error("system and chain must have matching beaddef")
    s.global_params == c.global_params || error("system and chain must have matching global_params")
    Nbeads= size(c.init_pos)[1]
    if isnothing(s.init_pos)
        init_pos= c.init_pos
        init_vel= c.init_vel
        perbead_params= c.perbead_params
    else
        init_pos= [s.init_pos; c.init_pos]
        init_vel= [s.init_vel; c.init_vel]
        perbead_params= [s.perbead_params; c.perbead_params]
    end
    chainbounds= (s.chainbounds...,s.chainbounds[end]+Nbeads)
    news= ChainSystem(s.starttime,
                    s.sp_pes,s.sp_forces,s.sp_interactions,s.sp_params,
                    s.beaddef,
                    perbead_params,s.global_params,init_pos, init_vel,chainbounds)
    #check energy and forces for issues
    force_pe(news,news.init_pos,news.starttime)
    news
end

"""
Return a new ChainSystem with added force
"""
function append(s::ChainSystem, f::SpecificForce)
    #error checking
    s.beaddef.VNdims == f.VNdims || error("system and force must have matching beaddef")
    pes= (s.sp_pes..., f.pe)
    forces= (s.sp_forces..., f.force)
    #convert interactions to full bead id, instead of tuple.
    beadids= map(i->map(ids->tobeadid(ids...,s.chainbounds), i), f.interactions)
    interactions= (s.sp_interactions..., beadids)
    params= (s.sp_params..., f.params)
    news= ChainSystem(s.starttime,pes,forces,interactions,params,
                s.beaddef,
                s.perbead_params,s.global_params,s.init_pos, s.init_vel,s.chainbounds)
    #check energy and forces for issues
    force_pe(news,news.init_pos,news.starttime)
    news
end

"""
Return force, pe of the beads
"""
function force_pe(s::ChainSystem,pos,time)
    @assert size(s.init_pos) == size(pos)
    f= zero(pos)
    pe= 0.0
    for fi in 1:length(s.sp_forces)
        for interaction in s.sp_interactions[fi]
            rs= map(id->pos[id,:], interaction)
            fs= s.sp_forces[fi](rs,time,s.sp_params[fi])
            spe= s.sp_pes[fi](rs,time,s.sp_params[fi])
            all(isfinite.(vcat(fs..., spe))) || @warn "some forces or energies are not finite on specific force $(sf), interaction $(interaction)"
            all(abs.(vcat(fs..., spe)) .< 2.0^30) || @warn "some forces or energies are over 2^30 on specific force $(sf), interaction $(interaction))"
            for i in 1:length(interaction)
                f[interaction[i],:] .+= fs[i]
            end
            pe += spe
        end
    end
    Nbeads= size(pos)[1]
    #chain forces
    for beadid in 1:Nbeads
        previd = prev_beadid(beadid,s.chainbounds)
        nextid = next_beadid(beadid,s.chainbounds)
        p= pos[previd,:]
        b= pos[beadid,:]
        n= pos[nextid,:]
        pf, bf, nf = s.beaddef.chain_force(p,b,n,s.perbead_params[beadid], s.global_params)
        e= s.beaddef.chain_pe(p,b,n,s.perbead_params[beadid], s.global_params)
        all(isfinite.([pf;bf;nf;e])) || @warn "chain forces or energies are not finite on Bead number $(beadid)"
        all(abs.([pf;bf;nf;e]) .< 2.0^30) || @warn "chain forces or energies are over 2^30 on Bead number $(beadid)"
        pe +=e
        f[previd,:] .+= pf
        f[beadid,:] .+= bf
        f[nextid,:] .+= nf
    end
    f, pe
end

# """
# Return system stepper
# """
# function makeintegrator(kernel,s::ChainSystem)

# end


module UnivariateFunctions

using Dates
using SchumakerSpline
using GLM

# This includes timing functions for use of UnivariateFunctions with dates.
include("date_conversions.jl")
export years_between, years_from_global_base, period_length
# The abstrat type UnivariateFunction and all structs are implemented here.
# In addition we have operator reversals and some supporting functions.
include("0_structs_and_generic_reversals.jl")
export UnivariateFunction, Undefined_Function, PE_Function, Sum_Of_Functions, Piecewise_Function
export change_base_of_PE_Function, trim_piecewise_function, convert_to_linearly_rescale_inputs

# These all implement evaluation, calculus and operators for the main structs.
include("1_undefined_function.jl")
include("2_pe_functions.jl")
include("3_sum_of_functions.jl")
include("4_piecewise_functions.jl")
export evaluate, derivative, indefinite_integral

include("5_calculus.jl")
export  evaluate_integral, right_integral, left_integral

include("chebyshevs.jl")
export get_chevyshevs_up_to

include("6_splines_and_interpolation.jl")
export create_quadratic_spline, create_constant_interpolation_to_right
export create_constant_interpolation_to_left, create_linear_interpolation

include("7_regressions_and_approximation.jl")
export create_ols_approximation, create_chebyshev_approximation

# The following operators also have overloads
export +, -, *, /, ^
end

# In terms of UnivariateFunction types we have the following seniority:
# Undefined_Function
# PE_Function
# Sum_Of_Functions
# Piecewise_Function

# In particular a Sum_Of_Functions contains the first two types of functions but never
# another Sum_Of_Functions (although one can be used to construct a Sum_Of_Functions). If a sum
# of functions contains an undefined function then it contains nothing else.
# A Piecewise_Function can have pieces composed of the first three categories but not another
# Piecewise_Function

# In terms of multiple dispatch rules this means that we:
# * define operators of each type of UnivariateFunction with another of its own type.
# * For each we define operators (f + x) with scalars. The reverse operations (x + f) can be done
#    all at once except for powers, subtraction and division where it matters.
#  All operations between these types should be defined for every pair. The seniority rule means that
#  Where a * b (where a and b are types) is the same as b * a then we define b*a = *(b,a) and put the
# relevent code in b * a place. Undefined_Function should take prevedence over PE_Function and etc
# in the seniority list.


# This example solves the linear 1D rotating shallow water equations
# for the depth-integrated velocities u and v and for the
# surface elevation η in a fluid with constant depth h,
# gravitational acceleration g and Coriolis parameter f:
#
# ∂u/∂t - fv = -gh∂η/∂x - D,   (1)
# ∂v/∂t + fu = - D,            (2)
# ∂η/∂t + ∂u/∂x = - D.         (3)
#
# where D indicates a hyperviscous operator of
# the form ν*(-1)^nν*Laplacian^nν, where ν is the viscosity
# coefficient and nν is the order of the operator.
# The code time-steps Equations (1-3) simultaneously.
#
# The initial conditions prescribe a Gaussian bump larger
# than the deformation radius in a fluid at rest, with
# some random noise smaller than the deformation radius
# added on top. The system develops geostrophically-balanced
# jets around the Gaussian bump, while the smaller-scale noise
# propagates around as inertia-gravity waves.

using FourierFlows, PyPlot
using FFTW: rfft, irfft, fft, ifft
using LinearAlgebra: mul!, ldiv!
using Printf
using Random

function calcN!(N, sol, t, s, v, p, g)
  mul!(v.uh, g.rfftplan, v.u)
  mul!(v.vh, g.rfftplan, v.v)
  mul!(v.ηh, g.rfftplan, v.η)
  rhsu = p.f*v.vh - p.g*p.h*im*g.kr.*v.ηh
  rhsv = - p.f*v.uh
  rhsη = - im*g.kr.*v.uh
  N[:, 1] .= rhsu
  N[:, 2] .= rhsv
  N[:, 3] .= rhsη
  dealias!(N, g, g.kralias)
  nothing
end

struct Vars{Aphys, Atrans} <: AbstractVars
        u :: Aphys
        v :: Aphys
        η :: Aphys
       uh :: Atrans
       vh :: Atrans
       ηh :: Atrans
end

function Vars(::Dev, g::AbstractGrid{T}) where {Dev, T}
  @devzeros Dev T g.nx u v η
  @devzeros Dev Complex{T} (g.nkr) uh vh ηh
  Vars(u, v, η, uh, vh, ηh)
end

struct Params{T} <: AbstractParams
       ν :: T         # Hyperviscosity coefficient
      nν :: Int       # Order of the hyperviscous operator
      g  :: T         # Gravitational acceleration
      h  :: T         # Fluid depth
      f  :: T         # Coriolis parameter
end

function Equation(p, g::AbstractGrid{T}) where T
  L = zeros(T, g.nkr, 3)
  D = - p.ν*g.kr.^(2*p.nν)
  L[:, 1] .= D # u
  L[:, 2] .= D # v
  L[:, 3] .= D # η
  FourierFlows.Equation(L, calcN!, g)
end

"""
    updatevars!(prob)
Update the vars in v on the grid g with the solution in sol.
"""
function updatevars!(prob)
  v, g, sol, p = prob.vars, prob.grid, prob.sol, prob.params
  ldiv!(v.u, g.rfftplan, deepcopy(sol[:, 1]))
  ldiv!(v.v, g.rfftplan, deepcopy(sol[:, 2]))
  ldiv!(v.η, g.rfftplan, deepcopy(sol[:, 3]))
  nothing
end

function set_uvη!(u0, v0, η0, prob)
    g, sol, p, v = prob.grid, prob.sol, prob.params, prob.vars
    mul!(v.uh, g.rfftplan, u0)
    mul!(v.vh, g.rfftplan, v0)
    mul!(v.ηh, g.rfftplan, η0)
    sol[:, 1] = v.uh
    sol[:, 2] = v.vh
    sol[:, 3] = v.ηh
    updatevars!(prob)
end

function updateplot!(axs, prob)
  sca(axs[1])
  cla()
  plot(prob.grid.x, prob.vars.η)
  xlabel(L"$x$ [m]")
  ylabel(L"$\eta$ [m]")
  sca(axs[2])
  cla()
  plot(prob.grid.x, prob.vars.u/prob.params.h) # Plot depth-averaged u.
  xlabel(L"$x$ [m]")
  ylabel(L"$\bar{u}$ [m/s]")
  sca(axs[3])
  cla()
  plot(prob.grid.x, prob.vars.v/prob.params.h) # Plot depth-averaged v.
  xlabel(L"$x$ [m]")
  ylabel(L"$\bar{v}$ [m/s]")
  xlim(prob.grid.x[1], prob.grid.x[end])
  subplots_adjust(hspace=0)
  nothing
end

#---
nx = 256     # Resolution
Lx = 20000π  # Domain length
ν = 1e-32    # Viscosity
nν = 1       # Viscosity order (nν = 1 means Laplacian)
dt = 0.01
g = 10.0     # Gravitational acceleration
h = 200.0    # Fluid depth
f = 1e-2     # Coriolis parameter
dev = CPU()

nt = 100000

grd = OneDGrid(nx, Lx)
vars = Vars(dev, grd)
params = Params(ν, nν, g, h, f)
eqn = Equation(params, grd)
stepper = "FilteredRK4"
iplt = 100

# Set up problem.
prob = FourierFlows.Problem(eqn, stepper, dt, grd, vars, params, dev)

# Set up the initial conditions, in this case a large Gaussian bump
# with small-scale perturbations superimposed.
u0 = zeros(grd.nx)
v0 = zeros(grd.nx)
noise = deepcopy(u0)
noise = Random.randn!(noise)/10
s = 1e4
η0 = @. 3*exp(-grd.x.^2/s^2) + noise
set_uvη!(u0, v0, η0, prob)

# Set up figure.
fig, axs = subplots(nrows=3, figsize=(8, 4), sharex=true)
Ld = @sprintf "%d" sqrt(g*h)/f
suptitle(L"Deformation radius $L_d = \sqrt{gh}/f = $"*string(Ld)*" m")

sca(axs[1])
plot(prob.grid.x, prob.vars.η)
xlabel(L"$x$ [m]")
ylabel(L"$\eta$ [m]")

sca(axs[2])
plot(prob.grid.x, prob.vars.u/prob.params.h) # Plot depth-averaged u.
xlabel(L"$x$ [m]")
ylabel(L"$\bar{u}$ [m/s]")

sca(axs[3])
plot(prob.grid.x, prob.vars.v/prob.params.h) # Plot depth-averaged v.
xlabel(L"$x$ [m]")
ylabel(L"$\bar{v}$ [m/s]")
xlim(prob.grid.x[1], prob.grid.x[end])

subplots_adjust(hspace=0)
show()

# Time-step the problem.
for i=1:nt
    println(i)
    if rem(i, iplt)==0
        updateplot!(axs, prob)
    end
    stepforward!(prob)
    updatevars!(prob)
end


module ConstLab

using Tensors, Parameters
import NLsolve, ForwardDiff, DiffResults

# Stress and strain control
abstract type ControlType end
struct StrainControl <: ControlType
    f # f(t) -> strain
end
struct StressControl <: ControlType
    # TODO: Make this just a MixedControl instead,
    # no gain in having a separate type
    f # f(t) -> stress
end
struct MixedControl <: ControlType
    # TODO: improve this masking by e.g. making masks 4th order tensors
    strain
    strain_mask::Vector{Bool}
    stress
    stress_mask::Vector{Bool}
    function MixedControl(strain, strain_mask, stress, stress_mask)
        strain_mask = collect(strain_mask.data)
        stress_mask = collect(stress_mask.data)
        for (a, b) in zip(strain_mask, stress_mask)
            a ⊻ b || error("no can do")
        end
        return new(strain, strain_mask, stress, stress_mask)
    end
end

###############################################
###############################################
###############################################
###############################################
abstract type Material end

abstract type MaterialState end

####################
## Linear Elastic ##
####################
struct LinearElastic <: Material
    G::Float64
    K::Float64
end
struct LinearElasticState{D,N} <: MaterialState
    ε::SymmetricTensor{2,D,Float64,N}
    σ::SymmetricTensor{2,D,Float64,N}
end
function LinearElasticState{D}() where D
    z = zero(SymmetricTensor{2,D})
    return LinearElasticState(z, z)
end

function constitutive_driver(model::LinearElastic, ε::SymmetricTensor, state::LinearElasticState)
    G, K = model.G, model.K
    σ = 2G * dev(ε) + 3K * vol(ε)
    return LinearElasticState(ε, σ)
end

#############################
## Plastic Prototype Model ##
#############################
struct Plastic <: Material
    G::Float64
    K::Float64
    σ_y::Float64 # Yield stress
    H::Float64
    κ∞::Float64
    α∞::Float64
    r::Float64
end
struct PlasticState{D,N} <: MaterialState
    ε::SymmetricTensor{2,D,Float64,N}
    σ::SymmetricTensor{2,D,Float64,N}
    εₚ::SymmetricTensor{2,D,Float64,N}
    κ::Float64
    α::SymmetricTensor{2,D,Float64,N}
    μ::Float64
end
# Initial state
function PlasticState{D}() where D
    z = zero(SymmetricTensor{2,D})
    return PlasticState(z, z, z, 0.0, z, 0.0)
end

function elastic_tangent(model::Plastic)
    G, K = model.G, model.K
    # Use Lamé parameters
    λ = K - 2G/3
    μ = G
    δ(i,j) = i == j ? 1.0 : 0.0
    E = SymmetricTensor{4,3}() do i, j, k, l
        λ * δ(i,j) * δ(k,l) + μ * (δ(i,k) * δ(j,l) + δ(i,l) * δ(j,k))
    end
    return E
end

mise(s) = √(3/2) * norm(dev(s))

function constitutive_driver(model::Plastic, ε::SymmetricTensor, state::PlasticState;
        solver_params=Dict{Symbol,Any}())
    G, K, σ_y = model.G, model.K, model.σ_y

    # Assume elastic step
    Δε = ε - state.ε
    σ_tr = state.σ + 2G * dev(Δε) + K * tr(Δε) * one(Δε)
    Φ = mise(σ_tr - state.α) - σ_y - state.κ
    if Φ < 0
        # elastic
        σ = σ_tr
        εₚ = state.εₚ
        κ = state.κ
        α = state.α
        μ = state.μ
        dσdε = elastic_tangent(model)
    else
        # plastic; need to iterate
        σ, εₚ, κ, α, μ, dσdε = solve_local_problem(model, Δε, state; solver_params=solver_params)
    end
    return σ, dσdε, PlasticState(ε, σ, εₚ, κ, α, μ)
end

function pack_variables(σ, κ, α, μ::T) where T
    n = length(σ.data) + length(κ) + length(α.data) + length(μ)
    @assert n == 14 # hehe
    X = Vector{T}(undef, n)
    return pack_variables!(X, σ, κ, α, μ)
end

function pack_variables!(X, σ, κ, α, μ)
    tomandel!(X, σ)
    X[7] = κ
    tomandel!(X, α; offset=6+1)
    X[end] = μ
    return X
end

function extract_variables(X)
    N = 6
    σ = frommandel(SymmetricTensor{2,3}, X)
    κ = X[N+1]
    α = frommandel(SymmetricTensor{2,3}, X; offset=N+1)
    μ = X[end]
    return σ, κ, α, μ
end

function solve_local_problem(model::Plastic, Δε::SymmetricTensor, state::PlasticState{D,N}; solver_params=Dict{Symbol,Any}()) where {D,N}
    # extract material parameters
    @unpack G, K, σ_y, κ∞, α∞, r, H = model

    σₙ = state.σ
    εₚₙ = state.εₚ
    κₙ = state.κ
    αₙ = state.α
    μₙ = state.μ * 0 # ??
    σ_tr = σₙ + 2G * dev(Δε) + K * tr(Δε) * one(Δε)

    # Initial guess
    X0 = pack_variables(σ_tr, κₙ, αₙ, μₙ)

    # Residual function
    function residual!(R, X)
        # fill!(R, 0) # pack_variables! should overwrite
        # Unpack variables
        σ, κ, α, μ = extract_variables(X)
        σ_red_dev = dev(σ - α)

        # Compute residuals
        Rσ = σ - σ_tr + 3G*μ/mise(σ_red_dev) * σ_red_dev
        Rκ = κ - κₙ - r * H * μ * (1 - κ/κ∞)
        Rα = α - αₙ - (1-r) * H * μ * (σ_red_dev / mise(σ_red_dev) - α/α∞)
        RΦ = mise(σ_red_dev) - σ_y - κ

        # Scale stresses with σ_y to make the residual dimensionless
        Rσ /= σ_y
        Rκ /= σ_y
        Rα /= σ_y
        RΦ /= σ_y

        # Populate residual vector R
        pack_variables!(R, Rσ, Rκ, Rα, RΦ)
        return R
    end

    # Use NLsolve to solve non-linear system
    params = Dict{Symbol,Any}(:ftol=>1e-14, :show_trace=>false, :method => :newton)
    merge!(params, solver_params)
    R = similar(X0) # residual cache for constructing OnceDifferentiable (not used :( )
    cache = NLsolve.OnceDifferentiable(residual!, X0, R; autodiff=:forward)
    res = NLsolve.nlsolve(cache, X0; params...)
    NLsolve.converged(res) || error("local problem did not converge: $res")
    X = res.zero # Solution to the problem
    J = cache.DF # Extract Jacobian from last step

    # # Custom implementation of Newton's algorithm
    # X = X0 # zeros(14)
    # R = zeros(14)
    # cfg = ForwardDiff.JacobianConfig(residual!, R, X)
    # out = DiffResults.JacobianResult(X)
    # iters = 0
    # local J
    # while true; iters += 1
    #     # residual!(R, X)
    #     ForwardDiff.jacobian!(out, residual!, R, X, cfg)
    #     g = DiffResults.value(out)
    #     J = DiffResults.jacobian(out)
    #     if norm(g) < 1e-6
    #         @info "local problem converged" iters norm(g)
    #         break
    #     elseif iters > 20
    #         error("did not converge")
    #     end
    #     ΔX = J \ g
    #     X .-= ΔX
    # end

    # Extract variables from the solution X
    σ, κ, α, μ = extract_variables(X)
    # Extract ATS tensor from final Jacobian
    dRdε = elastic_tangent(model) # For fixed X
    Jinv = inv(J)
    Jinv_t = frommandel(SymmetricTensor{4,3}, Jinv; offset_i=0, offset_j=0)
    dσdε = (Jinv_t / σ_y) ⊡ dRdε # scale Jacobian with σ_y since residual scaled with σ_y

    # Compute plastic strain; dεₚdt = λ ν -> εₚ = εₚₙ + Δt λ ν = εₚₙ + μ ν
    ν = 3/2 * dev(σ - α) / mise(dev(σ - α))
    εₚ = εₚₙ + μ * ν

    return σ, εₚ, κ, α, μ, dσdε
end

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

"""
    integrate(model, time, s0)

Integrate the material response.
"""
function integrate(model::Material, ctrl::StrainControl, time, state::T;
        solver_params=Dict{Symbol,Any}()) where {T <: MaterialState}
    states = T[]
    for (i, t) in enumerate(time)
        ε = ctrl.f(t)
        σ, dσdε, state′ = constitutive_driver(model, ε, state; solver_params=solver_params)
        push!(states, state′)
        state = state′
    end
    return states
end

function integrate(model::Material, ctrl::MixedControl, time, state::T;
        solver_params=Dict{Symbol,Any}()) where {T <: MaterialState}
    states = T[]
    for (i, t) in enumerate(time)
        # Compute controlled strain
        ε_c = ctrl.strain(t)
        # Compute controlled stress
        σ_c = ctrl.stress(t)

        # Iterate to find the correct strain
        local state′
        iters = 0
        Δε = zero(state.ε) # initial guess
        # println("----- strain iterations -----")
        while true; iters += 1
            # Current guess
            ε = SymmetricTensor{2,3}(ε_c.data .* ctrl.strain_mask .+ (state.ε.data .+ Δε.data) .* ctrl.stress_mask)

            σ′, dσdε, state′ = constitutive_driver(model, ε, state; solver_params=solver_params)

            # Compute stress residual
            R = tovoigt(σ′ - σ_c)[ctrl.stress_mask]
            # println("norm(R) = ", norm(R))
            if norm(R)/model.σ_y < 1e-6 # TOL
                # println("strain iterations converged in $iters iterations")
                break
            elseif iters > 20
                error("strain iterations did not converge")
            end
            # Take a step in strain
            J = tovoigt(dσdε)[ctrl.stress_mask, ctrl.stress_mask]
            ΔΔε = - J \ R
            ΔΔε_a = zeros(6)
            ΔΔε_a[ctrl.stress_mask] .= ΔΔε
            ΔΔε_t = fromvoigt(SymmetricTensor{2,3}, ΔΔε_a)
            Δε += ΔΔε_t
        end
        # Done, store stuff
        push!(states, state′)
        state = state′
    end
    return states
end

end # module


abstract type FluidModel end

# NOTE: Should add a test for matching u0 with grid
mutable struct DiffusionModel{T<:Union{Function,Nothing}} <: FluidModel
    u0::Array{Float64,2}
    u_analytic::T

    grid::Grid
    parameters::Parameters

    Δt::Float64
    t::Array{Float64,1}

    function DiffusionModel(u0, u_analytic::T, grid, parameters) where T<:Union{Function,Nothing}
        Δt = parameters.σ*grid.Δx*grid.Δy/parameters.ν
        t = parameters.tspan[1]:Δt:parameters.tspan[2]
        return new{T}(u0, u_analytic, grid, parameters, Δt ,t)
    end
end
DiffusionModel(grid, parameters) = DiffusionModel(zeros(grid.ny, grid.nx), nothing, grid, parameters)
DiffusionModel(u0, grid, parameters) = DiffusionModel(u0, nothing, grid, parameters);


const gk_nodes = (
    (0.0,),
    (-1.732050807568877193e+00, 1.732050807568877193e+00),
    (-7.410953499945408529e-01, -2.861279576057058183e+00, -4.184956017672732287e+00, 7.410953499945408529e-01, 2.861279576057058183e+00, 4.184956017672732287e+00),
    (1.230423634027306035e+00, 2.596083115049202306e+00, 3.205333794499194422e+00, 5.187016039913656229e+00, 6.363394494336369611e+00, -1.230423634027306035e+00, -2.596083115049202306e+00, -3.205333794499194422e+00, -5.187016039913656229e+00, -6.363394494336369611e+00),
    (-3.035159417196289788e-01, -2.208370035755473193e+00, -3.673120807624038697e+00, -4.688944733396042786e+00, -5.738816857731674048e+00, -7.052048695931292244e+00, -7.809096669391978551e+00, -8.661390355138452790e+00, -9.691888484307325413e+00, 3.035159417196289788e-01, 2.208370035755473193e+00, 3.673120807624038697e+00, 4.688944733396042786e+00, 5.738816857731674048e+00, 7.052048695931292244e+00, 7.809096669391978551e+00, 8.661390355138452790e+00, 9.691888484307325413e+00),
    (-6.206215996851309757e-01, -9.041718401396522387e-01, -1.556920948898949275e+00, -1.874020592699235932e+00, -2.570161501011057670e+00, -3.588718867304420357e+00, -3.813668067407607332e+00, -4.519991606227620196e+00, -4.871869590277843720e+00, -5.474559779359854161e+00, -6.034481198100183796e+00, -6.707072145057108692e+00, -7.289858306075302607e+00, -7.460729234405580890e+00, -8.181394510092470540e+00, -8.612155515170002218e+00, -8.886347001304686799e+00, -9.295854667849322794e+00, -1.009182628500030354e+01, -1.049797509743559409e+01, -1.091154446366382835e+01, -1.133361676429119136e+01, -1.176537484480120277e+01, -1.220821525266299190e+01, -1.266385440024828846e+01, -1.313446784351190466e+01, -1.362290117407219192e+01, -1.413301728617947717e+01, -1.467031519464652867e+01, -1.524314486007864389e+01, -1.586543512793624977e+01, -1.656423516155702913e+01, -1.741041065963615253e+01, 6.206215996851309757e-01, 9.041718401396522387e-01, 1.556920948898949275e+00, 1.874020592699235932e+00, 2.570161501011057670e+00, 3.588718867304420357e+00, 3.813668067407607332e+00, 4.519991606227620196e+00, 4.871869590277843720e+00, 5.474559779359854161e+00, 6.034481198100183796e+00, 6.707072145057108692e+00, 7.289858306075302607e+00, 7.460729234405580890e+00, 8.181394510092470540e+00, 8.612155515170002218e+00, 8.886347001304686799e+00, 9.295854667849322794e+00, 1.009182628500030354e+01, 1.049797509743559409e+01, 1.091154446366382835e+01, 1.133361676429119136e+01, 1.176537484480120277e+01, 1.220821525266299190e+01, 1.266385440024828846e+01, 1.313446784351190466e+01, 1.362290117407219192e+01, 1.413301728617947717e+01, 1.467031519464652867e+01, 1.524314486007864389e+01, 1.586543512793624977e+01, 1.656423516155702913e+01, 1.741041065963615253e+01)
)
const gk_nodes_v = (
    [0.0],
    [-1.732050807568877193e+00, 1.732050807568877193e+00],
    [-7.410953499945408529e-01, -2.861279576057058183e+00, -4.184956017672732287e+00, 7.410953499945408529e-01, 2.861279576057058183e+00, 4.184956017672732287e+00],
    [1.230423634027306035e+00, 2.596083115049202306e+00, 3.205333794499194422e+00, 5.187016039913656229e+00, 6.363394494336369611e+00, -1.230423634027306035e+00, -2.596083115049202306e+00, -3.205333794499194422e+00, -5.187016039913656229e+00, -6.363394494336369611e+00],
    [-3.035159417196289788e-01, -2.208370035755473193e+00, -3.673120807624038697e+00, -4.688944733396042786e+00, -5.738816857731674048e+00, -7.052048695931292244e+00, -7.809096669391978551e+00, -8.661390355138452790e+00, -9.691888484307325413e+00, 3.035159417196289788e-01, 2.208370035755473193e+00, 3.673120807624038697e+00, 4.688944733396042786e+00, 5.738816857731674048e+00, 7.052048695931292244e+00, 7.809096669391978551e+00, 8.661390355138452790e+00, 9.691888484307325413e+00],
    [-6.206215996851309757e-01, -9.041718401396522387e-01, -1.556920948898949275e+00, -1.874020592699235932e+00, -2.570161501011057670e+00, -3.588718867304420357e+00, -3.813668067407607332e+00, -4.519991606227620196e+00, -4.871869590277843720e+00, -5.474559779359854161e+00, -6.034481198100183796e+00, -6.707072145057108692e+00, -7.289858306075302607e+00, -7.460729234405580890e+00, -8.181394510092470540e+00, -8.612155515170002218e+00, -8.886347001304686799e+00, -9.295854667849322794e+00, -1.009182628500030354e+01, -1.049797509743559409e+01, -1.091154446366382835e+01, -1.133361676429119136e+01, -1.176537484480120277e+01, -1.220821525266299190e+01, -1.266385440024828846e+01, -1.313446784351190466e+01, -1.362290117407219192e+01, -1.413301728617947717e+01, -1.467031519464652867e+01, -1.524314486007864389e+01, -1.586543512793624977e+01, -1.656423516155702913e+01, -1.741041065963615253e+01, 6.206215996851309757e-01, 9.041718401396522387e-01, 1.556920948898949275e+00, 1.874020592699235932e+00, 2.570161501011057670e+00, 3.588718867304420357e+00, 3.813668067407607332e+00, 4.519991606227620196e+00, 4.871869590277843720e+00, 5.474559779359854161e+00, 6.034481198100183796e+00, 6.707072145057108692e+00, 7.289858306075302607e+00, 7.460729234405580890e+00, 8.181394510092470540e+00, 8.612155515170002218e+00, 8.886347001304686799e+00, 9.295854667849322794e+00, 1.009182628500030354e+01, 1.049797509743559409e+01, 1.091154446366382835e+01, 1.133361676429119136e+01, 1.176537484480120277e+01, 1.220821525266299190e+01, 1.266385440024828846e+01, 1.313446784351190466e+01, 1.362290117407219192e+01, 1.413301728617947717e+01, 1.467031519464652867e+01, 1.524314486007864389e+01, 1.586543512793624977e+01, 1.656423516155702913e+01, 1.741041065963615253e+01]
)
const gk_nodes_0 = (0.0,)
const gk_nodes_1 = (-1.732050807568877193e+00, 1.732050807568877193e+00),
const gk_nodes_2 = (-7.410953499945408529e-01, -2.861279576057058183e+00, -4.184956017672732287e+00, 7.410953499945408529e-01, 2.861279576057058183e+00, 4.184956017672732287e+00)
const gk_nodes_3 = (1.230423634027306035e+00, 2.596083115049202306e+00, 3.205333794499194422e+00, 5.187016039913656229e+00, 6.363394494336369611e+00, -1.230423634027306035e+00, -2.596083115049202306e+00, -3.205333794499194422e+00, -5.187016039913656229e+00, -6.363394494336369611e+00)
const gk_nodes_4 =  (-3.035159417196289788e-01, -2.208370035755473193e+00, -3.673120807624038697e+00, -4.688944733396042786e+00, -5.738816857731674048e+00, -7.052048695931292244e+00, -7.809096669391978551e+00, -8.661390355138452790e+00, -9.691888484307325413e+00, 3.035159417196289788e-01, 2.208370035755473193e+00, 3.673120807624038697e+00, 4.688944733396042786e+00, 5.738816857731674048e+00, 7.052048695931292244e+00, 7.809096669391978551e+00, 8.661390355138452790e+00, 9.691888484307325413e+00)
const gk_nodes_5 = (-6.206215996851309757e-01, -9.041718401396522387e-01, -1.556920948898949275e+00, -1.874020592699235932e+00, -2.570161501011057670e+00, -3.588718867304420357e+00, -3.813668067407607332e+00, -4.519991606227620196e+00, -4.871869590277843720e+00, -5.474559779359854161e+00, -6.034481198100183796e+00, -6.707072145057108692e+00, -7.289858306075302607e+00, -7.460729234405580890e+00, -8.181394510092470540e+00, -8.612155515170002218e+00, -8.886347001304686799e+00, -9.295854667849322794e+00, -1.009182628500030354e+01, -1.049797509743559409e+01, -1.091154446366382835e+01, -1.133361676429119136e+01, -1.176537484480120277e+01, -1.220821525266299190e+01, -1.266385440024828846e+01, -1.313446784351190466e+01, -1.362290117407219192e+01, -1.413301728617947717e+01, -1.467031519464652867e+01, -1.524314486007864389e+01, -1.586543512793624977e+01, -1.656423516155702913e+01, -1.741041065963615253e+01, 6.206215996851309757e-01, 9.041718401396522387e-01, 1.556920948898949275e+00, 1.874020592699235932e+00, 2.570161501011057670e+00, 3.588718867304420357e+00, 3.813668067407607332e+00, 4.519991606227620196e+00, 4.871869590277843720e+00, 5.474559779359854161e+00, 6.034481198100183796e+00, 6.707072145057108692e+00, 7.289858306075302607e+00, 7.460729234405580890e+00, 8.181394510092470540e+00, 8.612155515170002218e+00, 8.886347001304686799e+00, 9.295854667849322794e+00, 1.009182628500030354e+01, 1.049797509743559409e+01, 1.091154446366382835e+01, 1.133361676429119136e+01, 1.176537484480120277e+01, 1.220821525266299190e+01, 1.266385440024828846e+01, 1.313446784351190466e+01, 1.362290117407219192e+01, 1.413301728617947717e+01, 1.467031519464652867e+01, 1.524314486007864389e+01, 1.586543512793624977e+01, 1.656423516155702913e+01, 1.741041065963615253e+01)

const gk_nodes_nonneg = (
    ((0.0,),),
    ( (0.0,), (1.732050807568877193e+00,) ),
    ( (0.0,), (1.732050807568877193e+00,), (7.410953499945408529e-01, 2.861279576057058183e+00, 4.184956017672732287e+00)),
    ( (0.0,), (1.732050807568877193e+00,), (7.410953499945408529e-01, 2.861279576057058183e+00, 4.184956017672732287e+00), (1.230423634027306035e+00,2.596083115049202306e+00, 3.205333794499194422e+00, 5.187016039913656229e+00, 6.363394494336369611e+00) ),
    ( (0.0,), (1.732050807568877193e+00,), (7.410953499945408529e-01, 2.861279576057058183e+00, 4.184956017672732287e+00), (1.230423634027306035e+00, 2.596083115049202306e+00, 3.205333794499194422e+00, 5.187016039913656229e+00, 6.363394494336369611e+00), (3.035159417196289788e-01, 2.208370035755473193e+00, 3.673120807624038697e+00, 4.688944733396042786e+00, 5.738816857731674048e+00, 7.052048695931292244e+00, 7.809096669391978551e+00, 8.661390355138452790e+00, 9.691888484307325413e+00) ),
    ( (0.0,), (1.732050807568877193e+00,), (7.410953499945408529e-01, 2.861279576057058183e+00, 4.184956017672732287e+00), (1.230423634027306035e+00, 2.596083115049202306e+00, 3.205333794499194422e+00, 5.187016039913656229e+00, 6.363394494336369611e+00), (3.035159417196289788e-01, 2.208370035755473193e+00, 3.673120807624038697e+00, 4.688944733396042786e+00, 5.738816857731674048e+00, 7.052048695931292244e+00, 7.809096669391978551e+00, 8.661390355138452790e+00, 9.691888484307325413e+00), (6.206215996851309757e-01, 9.041718401396522387e-01, 1.556920948898949275e+00, 1.874020592699235932e+00, 2.570161501011057670e+00, 3.588718867304420357e+00, 3.813668067407607332e+00, 4.519991606227620196e+00, 4.871869590277843720e+00, 5.474559779359854161e+00, 6.034481198100183796e+00, 6.707072145057108692e+00, 7.289858306075302607e+00, 7.460729234405580890e+00, 8.181394510092470540e+00, 8.612155515170002218e+00, 8.886347001304686799e+00, 9.295854667849322794e+00, 1.009182628500030354e+01, 1.049797509743559409e+01, 1.091154446366382835e+01, 1.133361676429119136e+01, 1.176537484480120277e+01, 1.220821525266299190e+01, 1.266385440024828846e+01, 1.313446784351190466e+01, 1.362290117407219192e+01, 1.413301728617947717e+01, 1.467031519464652867e+01, 1.524314486007864389e+01, 1.586543512793624977e+01, 1.656423516155702913e+01, 1.741041065963615253e+01) )
)

const gk_weights = (
    ((1.0,),),
    ((6.666666666666666297e-01,), (1.666666666666666574e-01,)),
    ((2.539682539682539542e-01,), (9.485094850948509737e-02,), (2.700743295779378661e-01, 7.996325470893532808e-03, 9.426945755651748365e-05)),
    ((3.034671998542058935e-01,), (6.409605468680758256e-02,), (2.083249916496088816e-01, -6.337224793373735440e-03, 6.012336945984782032e-05), (6.115173012524767471e-02,1.808523425479845181e-02, 2.884880436506751140e-03, 6.094808731468982961e-07, 8.629684602229885931e-10)),
    ((9.688245529284254809e-02,), (4.421164421898454711e-02,), (1.430993028968333969e-01, 1.768022258182954352e-03, 3.226518598373974471e-05), (9.372082806552459178e-02, 4.333498812272349072e-03, 9.858275829964838720e-04, 2.959075202307440677e-07, 4.225258439631110457e-10), (1.476557104026862410e-01, 1.551310987485935398e-02, 2.349403664659752112e-04, 3.309758709792034345e-06, 1.659544880938981831e-08, 4.566176367618685878e-12, 1.822427515491293625e-14, 1.877818931437289588e-17, 1.903035094013049694e-21)),
    ((1.205674705975654931e-01,), (-2.241690549232919737e-03,), (-2.510946060860437534e-02, 2.317853741494368004e-03, 2.157340473532454820e-05), (6.054609539583663547e-02, -6.868088867756604043e-03, 8.108079306253045517e-04, 1.756589052393765114e-07, 2.175702428714775840e-10), (1.166672226897644410e-01, 1.165794765237651055e-02, -2.555098045918444941e-04, 2.736249936842626933e-07, 7.699205480918628382e-09, 2.123779208924948428e-12, 8.335951137105458603e-15, -1.603802571689216247e-17, 6.351777619826335869e-22), (1.161449102169717779e-01, 9.049907887097373427e-02, 3.958048330358010514e-02, 2.349638855141306104e-02, 1.187429460595669506e-02, 4.283096239652469652e-04, 1.399993373051156139e-04, 4.675973952563834421e-06, 8.812581632544770710e-07, 3.318711099677051844e-08, 1.560543759860818455e-09, 2.360855272927444972e-11, 1.169869203766870029e-13, 1.013359032410814206e-13, 4.429742347393624813e-16, 3.262744522103358460e-17, 1.322870363778446040e-18, 2.717875933335900016e-20, 1.232134560467146984e-23, 1.914714214218525338e-25, 2.332204174318528461e-27, 2.179136932463778929e-29, 1.524177029773773434e-31, 7.750724758884266004e-34, 2.764556078556950881e-36, 6.610498069228468602e-39, 9.992242712355998473e-42, 8.823151988713996391e-45, 4.072624717633741158e-48, 8.321030089138647840e-52, 5.728650998340608745e-56, 7.912918315856256558e-61, 5.881010699300556648e-67))
)

const weighted_gk_weights = (
    ((1.0,),),
    ((6.666666666666666297000000000000000000000000000000000000000000000000000000000002e-01,), (7.469481783896772988457527699794061672980745700022560408596162118340808426721962e-01,7.469481783896772988457527699794061672980745700022560408596162118340808426721962e-01)),
    ((2.539682539682539541999999999999999999999999999999999999999999999999999999999988e-01,), (4.250924592461578472759876805552536329599586523523682709259908670256577684243761e-01,4.250924592461578472759876805552536329599586523523682709259908670256577684243761e-01), (3.554229094726574982860442760983532555588167824313428698590967717008450205251302e-01, 4.793555157593584816398270172093156025646331160551993334837656250436617254414600e-01, 5.990412180534271700286076322376124144529767305328254509252998705048447247649803e-01,
    3.554229094726574982860442760983532555588167824313428698590967717008450205251302e-01, 4.793555157593584816398270172093156025646331160551993334837656250436617254414600e-01, 5.990412180534271700286076322376124144529767305328254509252998705048447247649803e-01)),
    ((3.034671998542058935000000000000000000000000000000000000000000000000000000000000e-01,), (2.872585877416563872923859174001655426169576823366256035985979317280372371506416e-01,
    2.872585877416563872923859174001655426169576823366256035985979317280372371506416e-01,), (2.741596165903048231254842627399873640947342306275468152337277728975795964203473e-01, -3.798974504437236345476663821676787033819236760209767324667376204925070214866375e-01, 3.82057746042619577466460231832021988947777427472900800439360113662814069216201e-01,
    2.741596165903048231254842627399873640947342306275468152337277728975795964203473e-01, -3.798974504437236345476663821676787033819236760209767324667376204925070214866375e-01, 3.82057746042619577466460231832021988947777427472900800439360113662814069216201e-01), (1.303638405254942398659921882662298597171952334211712030622447379181678924224902e-01, 5.257993039121150066942942760063309508915527286203890571040318547131226910007951e-01, 4.910597559959024396827368462951747108311140147810245954105871604633544199013856e-01, 4.239708523478058208142833543104371604410325009662179294092721824359687679500878e-01, 5.356639749505990342664174794050651607839649596753828921013245685339949128066048e-01,
    1.303638405254942398659921882662298597171952334211712030622447379181678924224902e-01, 5.257993039121150066942942760063309508915527286203890571040318547131226910007951e-01, 4.910597559959024396827368462951747108311140147810245954105871604633544199013856e-01, 4.239708523478058208142833543104371604410325009662179294092721824359687679500878e-01, 5.356639749505990342664174794050651607839649596753828921013245685339949128066048e-01)),
    ((9.688245529284254808999999999999999999999999999999999999999999999999999999999973e-02,), (1.981428426778980984412600814819064066197916004754116602341340067537383921399099e-01,
    1.981428426778980984412600814819064066197916004754116602341340067537383921399099e-01), (1.883213804828653468614944070827636503374772062253660876775346574368894296123359e-01, 1.059875844886804440092869110403572114614088909440086983015607835246120812498014e-01, 2.050311608171923251432647743399691597188646613280832007288539163124739211851223e-01,
    1.883213804828653468614944070827636503374772062253660876775346574368894296123359e-01, 1.059875844886804440092869110403572114614088909440086983015607835246120812498014e-01, 2.050311608171923251432647743399691597188646613280832007288539163124739211851223e-01), (1.997949536150074469472153344857304638452646615078951601015156424405452101031627e-01, 1.259895573866964685133959071955945228524379013064007907876920512461750179612389e-01, 1.67806002021515926781812697325227273981706539769413320575134754468269761718468e-01, 2.058410183088985338488574359666944058810790080238993669350031543082753239694951e-01, 2.622713152670181245847792066189062454897386919633485046840126633035438008876432e-01,
    1.997949536150074469472153344857304638452646615078951601015156424405452101031627e-01, 1.259895573866964685133959071955945228524379013064007907876920512461750179612389e-01, 1.67806002021515926781812697325227273981706539769413320575134754468269761718468e-01, 2.058410183088985338488574359666944058810790080238993669350031543082753239694951e-01, 2.622713152670181245847792066189062454897386919633485046840126633035438008876432e-01), (1.546159416265305613521062647988422327207440056433022066888637894519247591069617e-01, 1.777067237741947016011107256772430707943478489463044630983584233106965784410551e-01, 1.998335313964808739470910767885214844575959156383317195911069335120796374817492e-01, 1.968065697283842595393871134022617467268891580336636978901510751590663220829677e-01, 2.352449695913756290871153743077223673253485830732344148566236444253501744805265e-01, 2.874663303616266215185713034517514646718100831157478554314080989203956121474342e-01, 3.182154869178219322525709717261600085800217046015508192285268999262191679309989e-01, 3.66412995642927566567209315391227998614780120255270692272877691761267033639734e-01, 4.749785056705039548294688635784742503358477008080817452603676938307273765655888e-01,
    1.546159416265305613521062647988422327207440056433022066888637894519247591069617e-01, 1.777067237741947016011107256772430707943478489463044630983584233106965784410551e-01, 1.998335313964808739470910767885214844575959156383317195911069335120796374817492e-01, 1.968065697283842595393871134022617467268891580336636978901510751590663220829677e-01, 2.352449695913756290871153743077223673253485830732344148566236444253501744805265e-01, 2.874663303616266215185713034517514646718100831157478554314080989203956121474342e-01, 3.182154869178219322525709717261600085800217046015508192285268999262191679309989e-01, 3.66412995642927566567209315391227998614780120255270692272877691761267033639734e-01, 4.749785056705039548294688635784742503358477008080817452603676938307273765655888e-01)),
    ((1.205674705975654931000000000000000000000000000000000000000000000000000000000000e-01,), (-1.004656003357730823806188974619091124793125834446191621452480233775848578618848e-02,
    -1.004656003357730823806188974619091124793125834446191621452480233775848578618848e-02), (-3.304452355300148100907701108660353984545851091544121345792423435861862467492072e-02, 1.389483181685245561998489162800461481863219504281205217919283409423287612744794e-01, 1.370895620403929221796390803646287110031864827723960976487722418310067161769682e-01,
    -3.304452355300148100907701108660353984545851091544121345792423435861862467492072e-02, 1.389483181685245561998489162800461481863219504281205217919283409423287612744794e-01, 1.370895620403929221796390803646287110031864827723960976487722418310067161769682e-01), (1.290727426429006555492847049521242122808175102340887965026895848184566932833196e-01, -1.996787155198069833761048469679109478999974584765022381562309378357435651854228e-01, 1.380144353762277285999381191666272523336755548897420894976987031710802154905806e-01, 1.221929334587518376259432748909909467881910101233360640314212238681512645094766e-01, 1.350507538796819284118108665412593848348812381302706974707395964723043592556429e-01,
    1.290727426429006555492847049521242122808175102340887965026895848184566932833196e-01, -1.996787155198069833761048469679109478999974584765022381562309378357435651854228e-01, 1.380144353762277285999381191666272523336755548897420894976987031710802154905806e-01, 1.221929334587518376259432748909909467881910101233360640314212238681512645094766e-01, 1.350507538796819284118108665412593848348812381302706974707395964723043592556429e-01), (1.221667109516808190370518250428476817502289488961063083653087292766837791614186e-01, 1.335448340111544981514326484249212091804964172890004698602494327928524979978424e-01, -2.173293049894328541785639712838913063596493364318502740867021633896889137893188e-01, 1.627042969616786323017137575123847206390897111941153829472217340886271988531726e-02, 1.091383173808271059912672611652228508086363209626270939980457402609529582966774e-01, 1.337037745667200228518061287391775756922632633663635424192783589137993003687799e-01, 1.455546916115390192433400169433867669456458997723005332002505712744662073949346e-01, -3.129450315333034711164409656696647940753668072869003756323156170457042397932826e-01, 1.585340097882490853220738371167806400656320914148613019992326744838319491407595e-01,
    1.221667109516808190370518250428476817502289488961063083653087292766837791614186e-01, 1.335448340111544981514326484249212091804964172890004698602494327928524979978424e-01, -2.173293049894328541785639712838913063596493364318502740867021633896889137893188e-01, 1.627042969616786323017137575123847206390897111941153829472217340886271988531726e-02, 1.091383173808271059912672611652228508086363209626270939980457402609529582966774e-01, 1.337037745667200228518061287391775756922632633663635424192783589137993003687799e-01, 1.455546916115390192433400169433867669456458997723005332002505712744662073949346e-01, -3.129450315333034711164409656696647940753668072869003756323156170457042397932826e-01, 1.585340097882490853220738371167806400656320914148613019992326744838319491407595e-01), (1.408117965428582459162876149363881074540291961667053616321683414171503557489225e-01, 1.361970903861409165299347684909556849651325954647563203911887203853041023300648e-01, 1.329984629599679201317532559224692811885529412841919573627476858474076436367604e-01, 1.360229092078236952036773310956155750965076532404155400226276895185758575773446e-01, 3.228671754698710762567279580670329963945941193375398234348914314055081946042006e-01, 2.681489553918453044734559581543139773799622567549983705451696642331003349397749e-01, 2.015252077912545915669213003253490099523789818885819528195164688838272053088035e-01, 1.277204963843742976731996322288040800440204600020807221200522741183918951669275e-01, 1.256369893314551374852865986062839874863648830088856658669781405776860179840158e-01, 1.069170615807775867674670213187129633032953515717362686527691924792778806061541e-01, 1.260911971045680313452531871444771489334082184216232741175732749569181962079643e-01, 1.384835564679941519175122395707968073610522851416296138819888353652825409237885e-01, 4.053063396124052977373168417958918070128745886215735357114089793458949085693095e-02, 1.237990902870524369339555228492741317324406532190547262207679352134549377006608e-01, 1.517665185116327857270968738247869921185266823756355447386381896991320724004838e-01, 4.161262608892393519996858911156575973478844323821317971560465194785734821058607e-01, 1.857885153653686170918007032119532731433290633576380885676252281072935041150899e-01, 1.579630413832698922610809900762769849144924553399148831499524839581697316463859e-01, 1.606977230566520625746424052702404009903917533249360598570248398470460477209748e-01, 1.634385213505341330257697809858487597950195971180970507490215809799994145417857e-01, 1.666128532912926033083987912253689356198430144454419623138038929938517981765681e-01, 1.702300841125100847403377929828536379058458801150613409786168681599088845609372e-01, 1.743552051527007763833944392826094663092892128050440505451002005909832494783034e-01, 1.790935510945458777310090095483319171925624763449518313998782327281702016680291e-01, 1.845971971113272278601583938297960255043302296508511695418342101046108124115428e-01, 1.910850124205492485686214821819739358069442477378959613913593635449865408945723e-01, 1.988813937429899161582019334591191093761050728102145147551721737166719001341179e-01, 2.084906246429884317899166651270953944098724445617473229600911586848007892152892e-01, 2.207503567962314258656209780836950928147267874812765663538672701189731196089008e-01, 2.371870598593791815084402714470599016010287742002505126151560513044627873269844e-01, 2.60991031806506170012377255229863125256172232871574431654511691260342019258654e-01, 3.005170538128573338790480965022348051757581084764165815972783299215065489942171e-01, 3.905212369293002235972355867990859074521287143810981639577764587251872598475516e-01,
    1.408117965428582459162876149363881074540291961667053616321683414171503557489225e-01, 1.361970903861409165299347684909556849651325954647563203911887203853041023300648e-01, 1.329984629599679201317532559224692811885529412841919573627476858474076436367604e-01, 1.360229092078236952036773310956155750965076532404155400226276895185758575773446e-01, 3.228671754698710762567279580670329963945941193375398234348914314055081946042006e-01, 2.681489553918453044734559581543139773799622567549983705451696642331003349397749e-01, 2.015252077912545915669213003253490099523789818885819528195164688838272053088035e-01, 1.277204963843742976731996322288040800440204600020807221200522741183918951669275e-01, 1.256369893314551374852865986062839874863648830088856658669781405776860179840158e-01, 1.069170615807775867674670213187129633032953515717362686527691924792778806061541e-01, 1.260911971045680313452531871444771489334082184216232741175732749569181962079643e-01, 1.384835564679941519175122395707968073610522851416296138819888353652825409237885e-01, 4.053063396124052977373168417958918070128745886215735357114089793458949085693095e-02, 1.237990902870524369339555228492741317324406532190547262207679352134549377006608e-01, 1.517665185116327857270968738247869921185266823756355447386381896991320724004838e-01, 4.161262608892393519996858911156575973478844323821317971560465194785734821058607e-01, 1.857885153653686170918007032119532731433290633576380885676252281072935041150899e-01, 1.579630413832698922610809900762769849144924553399148831499524839581697316463859e-01, 1.606977230566520625746424052702404009903917533249360598570248398470460477209748e-01, 1.634385213505341330257697809858487597950195971180970507490215809799994145417857e-01, 1.666128532912926033083987912253689356198430144454419623138038929938517981765681e-01, 1.702300841125100847403377929828536379058458801150613409786168681599088845609372e-01, 1.743552051527007763833944392826094663092892128050440505451002005909832494783034e-01, 1.790935510945458777310090095483319171925624763449518313998782327281702016680291e-01, 1.845971971113272278601583938297960255043302296508511695418342101046108124115428e-01, 1.910850124205492485686214821819739358069442477378959613913593635449865408945723e-01, 1.988813937429899161582019334591191093761050728102145147551721737166719001341179e-01, 2.084906246429884317899166651270953944098724445617473229600911586848007892152892e-01, 2.207503567962314258656209780836950928147267874812765663538672701189731196089008e-01, 2.371870598593791815084402714470599016010287742002505126151560513044627873269844e-01, 2.60991031806506170012377255229863125256172232871574431654511691260342019258654e-01, 3.005170538128573338790480965022348051757581084764165815972783299215065489942171e-01, 3.905212369293002235972355867990859074521287143810981639577764587251872598475516e-01))
)

const weighted_gk_weights_nonneg = (
    ((1.0,),),
    ((6.666666666666666297000000000000000000000000000000000000000000000000000000000002e-01,), (7.469481783896772988457527699794061672980745700022560408596162118340808426721962e-01,)),
    ((2.539682539682539541999999999999999999999999999999999999999999999999999999999988e-01,), (4.250924592461578472759876805552536329599586523523682709259908670256577684243761e-01,), (3.554229094726574982860442760983532555588167824313428698590967717008450205251302e-01, 4.793555157593584816398270172093156025646331160551993334837656250436617254414600e-01, 5.990412180534271700286076322376124144529767305328254509252998705048447247649803e-01)),
    ((3.034671998542058935000000000000000000000000000000000000000000000000000000000000e-01,), (2.872585877416563872923859174001655426169576823366256035985979317280372371506416e-01,), (2.741596165903048231254842627399873640947342306275468152337277728975795964203473e-01, -3.798974504437236345476663821676787033819236760209767324667376204925070214866375e-01, 3.82057746042619577466460231832021988947777427472900800439360113662814069216201e-01), (1.303638405254942398659921882662298597171952334211712030622447379181678924224902e-01, 5.257993039121150066942942760063309508915527286203890571040318547131226910007951e-01, 4.910597559959024396827368462951747108311140147810245954105871604633544199013856e-01, 4.239708523478058208142833543104371604410325009662179294092721824359687679500878e-01, 5.356639749505990342664174794050651607839649596753828921013245685339949128066048e-01)),
    ((9.688245529284254808999999999999999999999999999999999999999999999999999999999973e-02,), (1.981428426778980984412600814819064066197916004754116602341340067537383921399099e-01,), (1.883213804828653468614944070827636503374772062253660876775346574368894296123359e-01, 1.059875844886804440092869110403572114614088909440086983015607835246120812498014e-01, 2.050311608171923251432647743399691597188646613280832007288539163124739211851223e-01), (1.997949536150074469472153344857304638452646615078951601015156424405452101031627e-01, 1.259895573866964685133959071955945228524379013064007907876920512461750179612389e-01, 1.67806002021515926781812697325227273981706539769413320575134754468269761718468e-01, 2.058410183088985338488574359666944058810790080238993669350031543082753239694951e-01, 2.622713152670181245847792066189062454897386919633485046840126633035438008876432e-01), (1.546159416265305613521062647988422327207440056433022066888637894519247591069617e-01, 1.777067237741947016011107256772430707943478489463044630983584233106965784410551e-01, 1.998335313964808739470910767885214844575959156383317195911069335120796374817492e-01, 1.968065697283842595393871134022617467268891580336636978901510751590663220829677e-01, 2.352449695913756290871153743077223673253485830732344148566236444253501744805265e-01, 2.874663303616266215185713034517514646718100831157478554314080989203956121474342e-01, 3.182154869178219322525709717261600085800217046015508192285268999262191679309989e-01, 3.66412995642927566567209315391227998614780120255270692272877691761267033639734e-01, 4.749785056705039548294688635784742503358477008080817452603676938307273765655888e-01)),
    ((1.205674705975654931000000000000000000000000000000000000000000000000000000000000e-01,), (-1.004656003357730823806188974619091124793125834446191621452480233775848578618848e-02,), (-3.304452355300148100907701108660353984545851091544121345792423435861862467492072e-02, 1.389483181685245561998489162800461481863219504281205217919283409423287612744794e-01, 1.370895620403929221796390803646287110031864827723960976487722418310067161769682e-01), (1.290727426429006555492847049521242122808175102340887965026895848184566932833196e-01, -1.996787155198069833761048469679109478999974584765022381562309378357435651854228e-01, 1.380144353762277285999381191666272523336755548897420894976987031710802154905806e-01, 1.221929334587518376259432748909909467881910101233360640314212238681512645094766e-01, 1.350507538796819284118108665412593848348812381302706974707395964723043592556429e-01), (1.221667109516808190370518250428476817502289488961063083653087292766837791614186e-01, 1.335448340111544981514326484249212091804964172890004698602494327928524979978424e-01, -2.173293049894328541785639712838913063596493364318502740867021633896889137893188e-01, 1.627042969616786323017137575123847206390897111941153829472217340886271988531726e-02, 1.091383173808271059912672611652228508086363209626270939980457402609529582966774e-01, 1.337037745667200228518061287391775756922632633663635424192783589137993003687799e-01, 1.455546916115390192433400169433867669456458997723005332002505712744662073949346e-01, -3.129450315333034711164409656696647940753668072869003756323156170457042397932826e-01, 1.585340097882490853220738371167806400656320914148613019992326744838319491407595e-01), (1.408117965428582459162876149363881074540291961667053616321683414171503557489225e-01, 1.361970903861409165299347684909556849651325954647563203911887203853041023300648e-01, 1.329984629599679201317532559224692811885529412841919573627476858474076436367604e-01, 1.360229092078236952036773310956155750965076532404155400226276895185758575773446e-01, 3.228671754698710762567279580670329963945941193375398234348914314055081946042006e-01, 2.681489553918453044734559581543139773799622567549983705451696642331003349397749e-01, 2.015252077912545915669213003253490099523789818885819528195164688838272053088035e-01, 1.277204963843742976731996322288040800440204600020807221200522741183918951669275e-01, 1.256369893314551374852865986062839874863648830088856658669781405776860179840158e-01, 1.069170615807775867674670213187129633032953515717362686527691924792778806061541e-01, 1.260911971045680313452531871444771489334082184216232741175732749569181962079643e-01, 1.384835564679941519175122395707968073610522851416296138819888353652825409237885e-01, 4.053063396124052977373168417958918070128745886215735357114089793458949085693095e-02, 1.237990902870524369339555228492741317324406532190547262207679352134549377006608e-01, 1.517665185116327857270968738247869921185266823756355447386381896991320724004838e-01, 4.161262608892393519996858911156575973478844323821317971560465194785734821058607e-01, 1.857885153653686170918007032119532731433290633576380885676252281072935041150899e-01, 1.579630413832698922610809900762769849144924553399148831499524839581697316463859e-01, 1.606977230566520625746424052702404009903917533249360598570248398470460477209748e-01, 1.634385213505341330257697809858487597950195971180970507490215809799994145417857e-01, 1.666128532912926033083987912253689356198430144454419623138038929938517981765681e-01, 1.702300841125100847403377929828536379058458801150613409786168681599088845609372e-01, 1.743552051527007763833944392826094663092892128050440505451002005909832494783034e-01, 1.790935510945458777310090095483319171925624763449518313998782327281702016680291e-01, 1.845971971113272278601583938297960255043302296508511695418342101046108124115428e-01, 1.910850124205492485686214821819739358069442477378959613913593635449865408945723e-01, 1.988813937429899161582019334591191093761050728102145147551721737166719001341179e-01, 2.084906246429884317899166651270953944098724445617473229600911586848007892152892e-01, 2.207503567962314258656209780836950928147267874812765663538672701189731196089008e-01, 2.371870598593791815084402714470599016010287742002505126151560513044627873269844e-01, 2.60991031806506170012377255229863125256172232871574431654511691260342019258654e-01, 3.005170538128573338790480965022348051757581084764165815972783299215065489942171e-01, 3.905212369293002235972355867990859074521287143810981639577764587251872598475516e-01))
)
const rule_lengths = (1,2,6,10,18,66)
const rules_total = (1,3,9,19,37,103)
const path = "/home/chris/Documents/progwork/julia/hermite/rules_nw"

function create_filename(rule)
    file = path * "/rule_"
    for i ∈ 1:rule
        file *= string(rules_list[i]) * "_"
    end
    file *= "qr_0.csv"
    file
end

"""
Creates and outputs sorted symmetric nested rules
so that we
"""
function print_rules(rule)
    out = Vector{Matrix{String}}(uninitialized, rule)
    floats = Vector{Vector{Float64}}(uninitialized, rule)
    for i ∈ rule:-1:1
        file = create_filename(i)
        temp_nodes = round.(@view(readdlm( file, ',')[1,:]),10)
        non_negative_temp_nodes = temp_nodes .>= 0
        out[i] = readdlm( file, ',', String)[[1,3],non_negative_temp_nodes]
        floats[i] = temp_nodes[non_negative_temp_nodes]
        if i < rule
            fi = floats[i]
            fip1 = floats[i+1]
            fadditions = fip1[[ fip1[k] ∉ fi for k ∈ eachindex(fip1) ]]
            n_ip1 = length(fip1)
        end
        for j ∈ i+1:rule
            fj = floats[j]
            inds = 1:n_ip1
            new_inds = [ fj[k] ∈ fadditions for k ∈ 1:n_ip1 ]
            sortperm_new_inds = sortperm(fj[inds][new_inds])
    #        if i > 1
                sortperm_unsorted = new_inds .== false
    #            @show sum(sortperm_unsorted)
                ind_sort = vcat( inds[sortperm_unsorted], (inds[new_inds])[sortperm_new_inds], n_ip1+1:length(fj) )
     #       else
     #           ind_sort = vcat( inds[new_inds][sortperm_new_inds], n_ip1+1:length(fj) )
     #       end
            fjsorted = fj[ind_sort]
            # @show (i,j)
            # @show fj
            # @show ind_sort
            # @show fjsorted
            # @show (length(fj), length(ind_sort))
            # println("\n")
            floats[j] = fjsorted
            out[j] = out[j][:,ind_sort]
        end
    end
    out
end

function print_string_mat(x)
    reduce((a,b) -> a * "," * b, x)
end
print_weighted_weights(A) = print_weighted_weights(A[1,:], A[2,:])
function print_weighted_weights(nodes, weights)
    out = ""
    for i in eachindex(nodes)
        n = BigFloat(strip(nodes[i]))
        w = BigFloat(strip(weights[i]))
        out *= string(exp(n^2/2)*w) * ", "
    end
    out
end

#using Reduce, SpecialFunctions
#Reduce.Rational(false);

gen_diff(v, i) = v > 0 ? :( $(Symbol('`' + i, :_, v)) - $(Symbol('`' + i, :_, v-1)) ) : Symbol('`' + i, :_, v)
function tupGridd(x::NTuple{N,T}) where {N,T<:Integer}
    Expr(:call, :*, ntuple(i -> gen_diff(x[i], i), Val{N}())... )
end
@generated function max_degree_poly(::Val{N}, l::Int, T = Int) where N
    quote
        j_0 = l
        s_0 = 0
        ind = 0
        out = Vector{NTuple{N,T}}(undef, binomial($N+l,$N))
        @nloops $N i p -> begin
            0:j_{$N-p}
        end p -> begin
            s_{$N-p+1} = s_{$N-p} + i_p
            j_{$N-p+1} = l - s_{$N-p+1}
        end begin
            ind += 1
            out[ind] = @ntuple $N j -> T(i_{j})
        end
        out
    end
end
function weights(::Val{N}, l::Int) where N
    p1 = max_degree_poly(Val{N}(), l)
    s1 = process_weights_no_deriv(p1)
    rcall(s1, :expand)::Expr
end
function process_weights_no_deriv(x::Vector{NTuple{N,Int}}) where N
    if length(x) == 1
        return tupGridd(x[1])
    end
    expr = Expr(:call, :+, tupGridd(x[1]), tupGridd(x[2]) )
    for i in 3:length(x)
        push!(expr.args, tupGridd(x[i]) )
    end
    expr
end

function inds(i::I, j::I, m::NTuple{N,I}, mcourse::NTuple{N,I}) where {N,I<:Integer}
    distbase = mcourse[i]
    distcurrent =  j-1 + distbase
    distcurrent, cld(hermite_weights[distcurrent],2) + m[i] - cld(length(hermite_weights[distbase]),2)
end

function weight(offstep::NTuple{N}, mcp1::NTuple{N}, mcourse::NTuple{N}, mfine::NTuple{N}, ::Type{T} = Float64) where {N,T}
    out = one(T)
    for k in 1:N
        i_k = offstep[k]
        mcp1_k = mcp1[k]
        mf_k = mfine[k]
        if i_k == 0
            out *= weighted_gk_weights[mcp1_k][mcp1_k][mf_k]
        else
            out *= weighted_gk_weights[i_k+mcp1_k][mcp1_k][mf_k] -  weighted_gk_weights[i_k+mcourse[k]][mcp1_k][mf_k]
        end
    end
    out
end
@generated function calc_weights(l_remaining::I, mcourse::NTuple{N,I}, mfine::NTuple{N,I}, ::Type{T} = Float64) where {N,I<:Integer,T}
    quote
        j_0 = l_remaining
        s_0 = 0
        out = zero($T)
        mcp1 = @ntuple $N i -> 1 + mcourse[i]
        @nloops $N i p -> begin
            0:j_{$N-p}
        end p -> begin
            s_{$N-p+1} = s_{$N-p} + i_p
            j_{$N-p+1} = l_remaining - s_{$N-p+1}
        end begin
            out += weight( (@ntuple $N k -> i_k), mcp1, mcourse, mfine, T )
        end
        out
    end
end

@generated function calc_weights2(l_remaining::I, mcourse::NTuple{N,I}, mfine::NTuple{N,I}, ::Type{T} = Float64) where {N,I<:Integer,T}
    quote
        j_0 = l_remaining
        s_0 = 0
        out = Vector{$T}(undef, 0)
        mcp1 = @ntuple $N i -> 1 + mcourse[i]
        @nloops $N i p -> begin
            0:j_{$N-p}
        end p -> begin
            s_{$N-p+1} = s_{$N-p} + i_p
            j_{$N-p+1} = l_remaining - s_{$N-p+1}
        end begin
            push!(out, weight( (@ntuple $N k -> i_k), mcp1, mcourse, mfine, T ))
        end
        out
    end
end


using JuMP
using LightGraphs, SimpleWeightedGraphs
using Random

lazy_cec_constraints = 0

function cycle_elimination_constraints!(
    model,
    graph::SimpleWeightedGraph,
    k::Int)

    global lazy_cec_constraints = 0
    @constraint(model, sum(variable_by_name(model, "y[1,$j]") for j in 2:nv(graph)) == 1) # Constraint 32
    @constraint(model, sum(variable_by_name(model, "y[$j,1]") for j in 2:nv(graph)) >= 1) # Constraint 33
    @constraint(model, sum(variable_by_name(model, "y[$j,1]") for j in 2:nv(graph)) <= k-1) # Constraint 34
    for i in 2:nv(graph)
        @constraint(model, sum(variable_by_name(model, "y[$i,$j]")
                               for j in 1:nv(graph) if has_edge(graph, i, j))
                    >= sum(variable_by_name(model, "y[$j,$i]")
                           for j in 1:nv(graph) if has_edge(graph, j, i))) # Constraint 35

        @constraint(model, sum(variable_by_name(model, "y[$i,$j]")
                               for j in 1:nv(graph) if has_edge(graph, i, j))
                    <= sum(variable_by_name(model, "y[$j,$i]")
                           for j in 1:nv(graph) if has_edge(graph, j, i)) * (k-1)) # Constraint 36
        @constraint(model, sum(variable_by_name(model, "y[$j,$i]") for j in 1:nv(graph) if has_edge(graph, j, i)) <= 1) # Constraint 37
    end

    MOI.set(model,
            MOI.LazyConstraintCallback(),
            cec_lazy_clause_generator(model, graph, k))
end

function cec_lazy_clause_generator(model, graph::SimpleWeightedGraph, k::Int)
    es = collect(SimpleWeightedEdge, edges(graph))
    sort!(es, by=dst)
    sort!(es, by=src, alg=MergeSort)
    y::Dict{Tuple{Int,Int},VariableRef} = Dict{Tuple{Int,Int},VariableRef}()
    for e in es
        i::Int = src(e)
        j::Int = dst(e)
        y[i,j] = variable_by_name(model, "y[$i,$j]")
        y[j,i] = variable_by_name(model, "y[$j,$i]")
    end
    return function (cb_data)
        status = callback_node_status(cb_data, model)
        if status == MOI.CALLBACK_NODE_STATUS_FRACTIONAL
            @debug "Found fractional solution"
            return
        elseif status == MOI.CALLBACK_NODE_STATUS_INTEGER
            @debug "Found integer solution"
            # x_val::Dict{Int, } = Dict{Int, Float64}()
            y_val::Dict{Tuple{Int,Int},Int} = Dict{Tuple{Int,Int},Int}()
            for e in es
                i::Int = src(e)
                j::Int = dst(e)
                y_val[i,j] = round(callback_value(cb_data, y[i,j]))
                y_val[j,i] = round(callback_value(cb_data, y[j,i]))
            end
            cycle = find_cycle(graph, y_val,k)
            if !isnothing(cycle)
                cyclelength::Int = length(cycle)
                @debug "Detected cycle with length $cyclelength: $cycle."
                constraint = @build_constraint(sum(y[i,j]
                                                    for i in cycle
                                                        for j in cycle
                                                            if has_edge(graph, i, j))
                                               <= cyclelength - 1) # Constraint 38

                # for debugging:
                # arcs::Vector{String} = Vector{String}()
                # for i in cycle
                #    for j in cycle
                #        if has_edge(graph, i, j)
                #            push!(arcs, "y[$i,$j]")
                #        end
                #    end
                # end
                # @debug join(arcs, " + ") * " <= $(cyclelength - 1)"

                global lazy_cec_constraints += 1
                MOI.submit(model, MOI.LazyConstraint(cb_data), constraint)
            end
        else
            @assert status == MOI.CALLBACK_NODE_STATUS_UNKNOWN
            @debug "Found either fractional or integer solution"
        end
    end
end

function find_cycle(
    graph::SimpleWeightedGraph,
    y_val::Dict{Tuple{Int,Int},Int64},
    k::Int)::Union{Nothing,Set{Int}}
    visited::Set{Int} = Set{Int}()
    queue::Set{Int} = Set{Int}([1])
    while !isempty(queue)
        @assert !isempty(queue) "queue should not be empty."
        c::Int = pop!(queue)
        ns = neighbors(graph, c)
        for n::Int in ns
            if y_val[c, n] == 1 && n != 1
                push!(queue, n)
            end
        end
        push!(visited, c)
    end
    if length(visited) == k + 1
        return # no cycles
    end
    @assert length(visited) <= k "Solution should be a forest, but only has one tree."

    to_check::Set{Int} = Set{Int}(2:nv(graph))

    filter!(v -> !(v in visited), to_check)
    @assert length(to_check) == (nv(graph) - 1) - (length(visited) - 1) "Should have pruned $(length(visited) - 1) of $(nv(graph) - 1) items. However only $(length(to_check)) left."

    while !isempty(to_check)
        @assert isempty(queue) "queue should be empty."
        push!(queue, pop!(to_check))
        visited = Set{Int}()
        @assert isempty(visited) "visited should be empty."
        while !isempty(queue)
            @assert !isempty(queue) "queue should not be empty."
            c::Int = pop!(queue)
            # @debug "Visiting node $c"
            push!(visited, c)
            ns = neighbors(graph, c)
            for n::Int in ns
                if y_val[c, n] == 1 && n != 1
                    if n in visited
                        return visited
                    end
                    push!(queue, n)
                end
            end
        end
        filter!(v -> !(v in visited), to_check)
    end

    @warn "No cycle detected. Something is wrong with the implementation."

end


using BEAST

Δt, Nt = 1.0, 200
T0 = timebasisshiftedlagrange(Δt, Nt, 0)
T1 = timebasisshiftedlagrange(Δt, Nt, 1)
T2 = timebasisshiftedlagrange(Δt, Nt, 2)
T3 = timebasisshiftedlagrange(Δt, Nt, 3)
iT0 = integrate(T0)
iT1 = integrate(T1)
iT2 = integrate(T2)
iT3 = integrate(T3)
δ = timebasisdelta(Δt, Nt)

Id = BEAST.Identity()
Z = assemble(Id, δ, iT1)
Z = reshape(Z,1,1,length(Z))

taxis = range(0, step=Δt, length=Nt)
width, delay = 2.0, 6.0
f(t) = -2*(t-delay)/width^2*exp(-((t-delay)/width)^2)
F(t) = exp(-((t-delay)/width)^2)

b = F.(taxis)'
u = marchonintime(inv(Z[:,:,1]), Z, b, Nt)

plot(taxis, u[1,:])

taxis2 = range(-3Δt, 5Δt, length=200)
plot()
plot!(taxis2, iT3.(taxis2))

kmax = 2
cv = BEAST.ConvOp(
    reshape(Z[1,1,1:kmax],kmax,1,1),
    fill(1,(1,1)),
    fill(kmax,(1,1)),
    fill(1.0,(1,1)),
    Nt)

w, Q = BEAST.polyeig(cv)

u2 = marchonintime(inv(Z[:,:,1]), cv, b, Nt)


## Example from Bradlyn, Science, Eq. (S66)
# SG199
g3 = S"y+1,z,x+1"       # {3₁₁₁⁻|1,0,1}
g1 = S"x-1/2,-y+1/2,-z" # 2₁₀₀|-½,½,0
g2 = S"-x,y+1/2,-z-1/2" # {2₀₁₀|0,½,-½}

G3 = [0.0 0 1; 1 0 0; 0 1 0]
G1 = [-1.0 0 0; 0 -1 0; 0 0 1]
G2 = inv(G3)*G1*(G3)

lg = LittleGroup(199, KVec("0.5,0.5,0.5"), "P", [g1,g2,g3,])#g1, g2, g3]) #
lgir = LGIrrep(
    "P?",
    lg,
    complex.([G1,G2,G3,]),#, G2, G3]), #
    [zeros(3) for _ in 1:3],
    REAL,
    false
    )

H = kdotp(lgir)

## -----
# SG220

g3 = S"z,x,y"
g4 = S"-x+1/2,z+1,-y+1"
m = S"y+1/2,x+1/2,z+1/2"
G3 = -1.0im*[0 1 0; 0 0 -1; 1 0 0]
G4 = [0 0 1; 0 1 0; -1 0 0]*cispi(1/4)
M = [1 0 0; 0 0 1; 0 1 0]*cispi(3/4)

lg = LittleGroup(220, KVec("0.5,0.5,0.5"), "P", [g3, g4, m])
lgir = LGIrrep(
    "P?",
    lg,
    complex.([G3, G4, M]),
    [zeros(3) for _ in 1:3],
    REAL,
    false
    )
kdotp(lgir)


include("$(@__DIR__)/../src/FiniteGroups.jl")
using .FiniteGroups, LinearAlgebra, Test

g = pointgroup("Oh")
rreps = irreps(g, R=true)

rep = oplus(
    rreps[1],   # 1
    rreps[2],   # 2
    rreps[2],   # 3
    rreps[2],   # 4
    rreps[2],   # 5
    rreps[5],   # 6
    rreps[6],   # 7
    rreps[8],   # 8
    rreps[8],   # 9
    rreps[9],   # 10
    rreps[10]   # 11
)
ru = Symmetric(rand(size(rep[1])...)) |> eigvecs
rep = transform_rep(ru', rep, ru)
res = block_decomposition(rep, g)

@testset "Oh Decomp" begin
    @test length(res) == 11
    @test norm(transform_rep(res[1]', rep, res[1]) .- rreps[1]) < 1e-10
    for i = 2:5
        @test norm(transform_rep(res[i]', rep, res[i]) .- rreps[2]) < 1e-10
    end
    for i = 6:6
        @test norm(transform_rep(res[i]', rep, res[i]) .- rreps[5]) < 1e-10
    end
    for i = 7:7
        @test norm(transform_rep(res[i]', rep, res[i]) .- rreps[6]) < 1e-10
    end
    for i = 8:9
        @test norm(transform_rep(res[i]', rep, res[i]) .- rreps[8]) < 1e-10
    end
    for i = 10:10
        @test norm(transform_rep(res[i]', rep, res[i]) .- rreps[9]) < 1e-10
    end
    for i = 11:11
        @test norm(transform_rep(res[i]', rep, res[i]) .- rreps[10]) < 1e-10
    end
end


struct KNN
  kcoeff::Float64
  neighbors::Data{Weighted}
  weights::Vector{Vector{Float64}}
end
get_bestk(n, knn::KNN) = round(Int, knn.kcoeff * sqrt(n))

"Find the best rule for k on the given data `d`."
function find_k_coeff(d::Data{Weighted}, k_range=1:100)
  M = length(d.y_T)

  best_k = zeros(Int, M)
  for m = 1:M
    X_m = d.X_T[m]
    y_m = d.y_T[m]
    n_m = length(y_m)

    max_k = min(maximum(k_range), n_m - 1)
    kdtree = NearestNeighbors.KDTree(Matrix(X_m'))
    # Find the k+1 nearest neighbors to each point, sorted by distance
    # Each point has itself as the first neighbor, which we will ignore
    idxs, _ = NearestNeighbors.knn(kdtree, Matrix(X_m'), max_k + 1, true)

    # Get the outcomes for each of the neighbors
    # Ignore the first neighbor since it's just the point itself
    cum_means = Vector{Vector{Float64}}(undef, n_m)
    for i = 1:n_m
      outcomes = y_m[idxs[i][2:end]]
      cum_means[i] = cumsum(outcomes) ./ (1:max_k)
    end

    # Find the k that minimizes the MAE of kNN predictions
    best_mae = Inf
    for k in k_range
      k >= n_m && continue
      mae = 0.0
      for i = 1:n_m
        mae += abs(cum_means[i][k] - y_m[i])
      end
      mae /= n_m

      if mae < best_mae
        best_mae = mae
        best_k[m] = k
      end
    end
  end

  sqrt_n = sqrt.(length.(d.y_T))
  path = GLMNet.glmnet(reshape(sqrt_n, M, 1), best_k)
  k_coeff = path.betas[1, end]
end

function getoutcomes(newdata_norm::Data{Normalized}, knn::KNN)
  M = length(knn.neighbors.X_T)

  n = size(newdata_norm.X, 1)
  outcomes = Matrix{Float64}(undef, n, M)

  for m = 1:M
    # Weight all points according to candidate treatment
    current_weighted = weight_data(newdata_norm, knn.weights, m)

    # Get all relevant weighted training points for candidate treatment
    X_m = knn.neighbors.X_T[m]
    y_m = knn.neighbors.y_T[m]

    k = get_bestk(length(y_m), knn)
    k = max(min(k, length(y_m)), 1)

    # Get k nearest neighbors for each new point from the neighbor pool
    # We don't need the neighbors in sorted order
    kdtree = NearestNeighbors.KDTree(Matrix(X_m'))
    idxs, _ = NearestNeighbors.knn(kdtree, Matrix(current_weighted.X'), k,
                                   false)

    # Get the mean outcome for each point from the neighbors
    for i = 1:n
      outcome = 0.0
      idx = idxs[i]
      for j = 1:k
        outcome += y_m[idx[j]]
      end
      outcomes[i, m] = outcome / k
    end
  end

  outcomes
end

"Trains kNN model on training data `d`."
function train(d::Data{Normalized}, ::Val{:knn})
  d_weighted, weights = weight_data(d)

  k_coeff = find_k_coeff(d_weighted)
  KNN(k_coeff, d_weighted, weights)
end

"Estimates counterfactuals on the dataset `d` using kNN."
function getcounterfactuals(d::Data{Normalized}, impute_method::Val{:knn})
  knn = train_knn(d, impute_method)
  getoutcomes(d, knn)
end


# isometric projection
to_2d(x, y, z) = (y-x)*sqrt(3)/2, z-(x+y)/2

# number of blocks along one side of square base
const N = 20

# (x,y) of uppermost corner of each tower's base "diamond", from the
# back to the front (and left to right).
function roots()
    Task() do
        for sum in 0:2N-1
            x0 = min(sum, N-1)
            y0 = max(0, sum-N+1)
            for col in 0:(x0-y0)
                x = x0-col
                y = y0+col
                produce(x, y)
            end
        end
    end
end

# the "outline" of a tower (for filling)
function column_outline(x, y, z)
    move(to_2d(x+1, y, 0)...)
    line(to_2d(x+1, y+1, 0)...)
    line(to_2d(x, y+1, 0)...)
    line(to_2d(x, y+1, z)...)
    line(to_2d(x, y, z)...)
    line(to_2d(x+1, y, z)...)
    line(to_2d(x+1, y, 0)...)
end

# the outline extended with a "Y" (for stroking)
function column(x, y, z)
    column_outline(x, y, z)
    move(to_2d(x+1, y, z)...)
    line(to_2d(x+1, y+1, z)...)
    line(to_2d(x, y+1, z)...)
    move(to_2d(x+1, y+1, z)...)
    line(to_2d(x+1, y+1, 0)...)
end

# put it all together
with(PNG("towers.png", 300, 150), Axes(centred=true),
Ink("black"), Pen(0.03), Translate(0, 1), Scale(2.4/N)) do
    for (x, y) in roots()
        d = sqrt((x/N-1/2)^2 + (y/N-1/2)^2)
        z = max(0, cos(d * pi * 0.9) * N * 0.5 * rand())
        if z > 1
            paint(Ink("white")) do
                column_outline(x,y,z)
            end
            draw() do
                column(x,y,z)
            end
        end
    end
end

println("towers: ok")