text:
module ThinWalledBeam

using DiffEqOperators: CenteredDifference
using LinearAlgebra
using NLsolve

using ..Mesh

export discretize, define, solve, Model

Base.@kwdef mutable struct Model   #using Parameters.jl macro here to help assign some constants, leave others for later.

   Ix::Union{Array{Float64}, Nothing} = nothing
   Iy::Union{Array{Float64}, Nothing} = nothing
   Ixy::Union{Array{Float64}, Nothing} = nothing
   J::Union{Array{Float64}, Nothing} = nothing
   Cw::Union{Array{Float64}, Nothing} = nothing
   E::Union{Array{Float64}, Nothing} = nothing
   ν::Union{Array{Float64}, Nothing} = nothing
   G::Union{Array{Float64}, Nothing} = nothing
   ax::Union{Array{Float64}, Nothing} = nothing
   ay::Union{Array{Float64}, Nothing} = nothing
   ay_kx::Union{Array{Array{Float64}}, Nothing} = nothing
   kx::Union{Array{Array{Float64}}, Nothing} = nothing
   kϕ::Union{Array{Array{Float64}}, Nothing} = nothing

   qx::Union{Array{Float64}, Nothing} = nothing
   qy::Union{Array{Float64}, Nothing} = nothing

   end_boundary_conditions::Union{Array{Int64}, Nothing} = nothing
   supports::Union{Array{Float64}, Nothing} = nothing

   z::Union{Array{Float64}, Nothing} = nothing
   m::Union{Array{Int64}, Nothing} = nothing

   Kff::Union{Matrix{Float64}, Nothing} = nothing
   # Ksf::Union{Matrix{Float64}, Nothing} = nothing
   Ff::Union{Array{Float64}, Nothing} = nothing
   # Fs::Union{Array{Float64}, Nothing} = nothing

   free_dof::Union{Array{Int64}, Nothing} = nothing

   u::Union{Array{Float64}, Nothing} = nothing
   v::Union{Array{Float64}, Nothing} = nothing
   ϕ::Union{Array{Float64}, Nothing} = nothing

end

function calculate_derivative_operators(dz)

   #Define the number of nodes.
   num_nodes = length(dz) + 1

   #Add extra dz on each end for padding nodes used in the CenteredDifference function.
   dz = [dz[1]; dz; dz[end]]

   #Calculate the 4th derivative operator.
   nth_derivative = 4
   derivative_order = 2
   Azzzz = CenteredDifference(nth_derivative, derivative_order, dz, num_nodes)

   #Convert the operator to a matrix.
   Azzzz = Array(Azzzz)

   #Trim off the ghost nodes.
   Azzzz = Azzzz[:,2:end-1]

   #Calculate the 2nd derivative operator.
   nth_derivative = 2
   derivative_order = 2
   Azz = CenteredDifference(nth_derivative, derivative_order, dz, num_nodes)
   Azz = Array(Azz)
   Azz = Azz[:,2:end-1]

   return Azzzz, Azz

end

function calculate_boundary_stencils(bc_flag, h, nth_derivative)

   #Calculate boundary conditions stencils without ghost nodes using
   #Jorge M. Souza, "Boundary Conditions in the Finite Difference Method"
   #https://cimec.org.ar/ojs/index.php/mc/article/download/2662/2607

   taylor_coeffs =  [1; h; h^2/2; h^3/6; h^4/24]

   RHS = zeros(Float64, 5)
   #row1*[A;B;C;D;E]*u2
   #row2*[A;B;C;D;E]*u2'
   #row3*[A;B;C;D;E]*u2''
   #row4*[A;B;C;D;E]*u2'''
   #row5*[A;B;C;D;E]*u2''''

   if bc_flag == 1 #simply supported end

      LHS = [1 1 1 1 0
         -1 0 1 2 0
         1 0 1 4 2
         -1 0 1 8 -6
         1 0 1 16 12]

      RHS[nth_derivative + 1] = (1/taylor_coeffs[nth_derivative + 1])
      boundary_stencil = LHS \ RHS
      boundary_stencil = ((boundary_stencil[1:4]),(zeros(4)))  #since u''=0

   elseif bc_flag == 2 #fixed end

      LHS = [1 1 1 1 0
         -1 0 1 2 1
         1 0 1 4 -2
         -1 0 1 8 3
         1 0 1 16 -4]

      RHS[nth_derivative+1] = (1/taylor_coeffs[nth_derivative + 1])
      boundary_stencil = LHS\RHS
      boundary_stencil = ((boundary_stencil[1:4]),(zeros(4)))  #since u'=0

   elseif bc_flag == 3 #free end
                #u'' u'''
      LHS = [1 1 1  0   0
           0 1 2  0   0
           0 1 4  0.5 0
           0 1 8  0   6
           0 1 16 0  0]

      RHS[nth_derivative+1] = (1/taylor_coeffs[nth_derivative+1])
      boundary_stencil1 = LHS\RHS  #at free end
      boundary_stencil1 = boundary_stencil1[1:3]   #u''=u'''=0

      # use simply supported BC to find stencil at one node in from free end
      LHS = [1 1 1 1 0
         -1 0 1 2 0
         1 0 1 4 2
         -1 0 1 8 -6
         1 0 1 16 12]

      RHS[nth_derivative+1] = (1/taylor_coeffs[nth_derivative+1])
      boundary_stencil2 = LHS\RHS #at one node over from free end
      boundary_stencil2 = boundary_stencil2[1:4]
      boundary_stencil = ((boundary_stencil1), (boundary_stencil2))  #two stencils are calculated

   end

   return boundary_stencil

end

function apply_end_boundary_conditions(A, end_boundary_conditions, nth_derivative, dz)

   #Consider the left end boundary condition.  Swap out interior stencil for boundary stencil.
   h = dz[1]
   bc_flag = end_boundary_conditions[1]
   boundary_stencil = calculate_boundary_stencils(bc_flag, h, nth_derivative)

   A[1,:] .= 0.0
   A[2,:] .= 0.0

   if (bc_flag == 1) | (bc_flag == 2)   #make this cleaner, combine
      A[2,1:length(boundary_stencil[1])] = boundary_stencil[1]
   else
      A[1,1:length(boundary_stencil[1])] = boundary_stencil[1]
      A[2,1:length(boundary_stencil[2])] = boundary_stencil[2]
   end

   #Consider the right end boundary condition.
   h = dz[end]
   bc_flag = end_boundary_conditions[2]
   boundary_stencil = calculate_boundary_stencils(bc_flag,h,nth_derivative)

   A[end,:] .= 0.0
   A[end-1,:] .= 0.0

   if (bc_flag == 1) | (bc_flag == 2)
      A[end-1,(end-(length(boundary_stencil[1])-1)):end] = reverse(boundary_stencil[1])
   else
      A[end,end-(length(boundary_stencil[1])-1):end] = reverse(boundary_stencil[1])
      A[end-1,end-(length(boundary_stencil[2])-1):end] = reverse(boundary_stencil[2])
   end

   return A

end

function discretize(member_definitions)

      #Define discretization along beam (z) and assign a member type to each node (m).
      dz, z, m = Mesh.define_line_element(member_definitions)

      return z, m

end

function define(z, m, member_definitions, section_properties, material_properties, kx, kϕ, ay_kx, qx, qy, ax, ay, end_boundary_conditions, supports)

   #This is the user interface mapping.
   #L(1) dL(2) section_properties(3) material_properties(4)

   #Define properties at each node.
   dz = diff(z)
   Ix = Mesh.create_line_element_property_array(member_definitions, m, dz, section_properties, 3, 1)
   Iy = Mesh.create_line_element_property_array(member_definitions, m, dz, section_properties, 3, 2)
   Ixy = Mesh.create_line_element_property_array(member_definitions, m, dz, section_properties, 3, 3)
   J = Mesh.create_line_element_property_array(member_definitions, m, dz, section_properties, 3, 4)
   Cw = Mesh.create_line_element_property_array(member_definitions, m, dz, section_properties, 3, 5)

   E = Mesh.create_line_element_property_array(member_definitions, m, dz, material_properties, 4, 1)
   ν = Mesh.create_line_element_property_array(member_definitions, m, dz, material_properties, 4, 2)
   G = E./(2 .*(1 .+ ν))

   #Store model inputs in data structure.
   model = Model(Ix, Iy, Ixy, J, Cw, E, ν, G, ax, ay, ay_kx, kx, kϕ, qx, qy, end_boundary_conditions, supports, z, m, nothing, nothing, nothing, nothing, nothing, nothing)

   return model

end

function governing_equations(model)

   #Define the number of nodes.
   num_nodes = length(model.z)

   #Calculate the derivative operators.
   dz = diff(model.z)
   Azzzz,Azz = calculate_derivative_operators(dz)

   #Apply left and right end boundary condition stencils to derivative operators.
   nth_derivative = 4
   Azzzz = apply_end_boundary_conditions(Azzzz, model.end_boundary_conditions, nth_derivative, dz)

   nth_derivative = 2
   Azz = apply_end_boundary_conditions(Azz, model.end_boundary_conditions, nth_derivative, dz)

   #Build identity matrix for ODE operations.
   AI = Matrix(1.0I, num_nodes, num_nodes)

   #Build operator matrix that doesn't update with load.

   #LHS first.
   A11 = zeros(Float64, num_nodes,num_nodes)
   A12 = zeros(Float64, num_nodes,num_nodes)
   A13 = zeros(Float64, num_nodes,num_nodes)
   A21 = zeros(Float64, num_nodes,num_nodes)
   A22 = zeros(Float64, num_nodes,num_nodes)
   A23 = zeros(Float64, num_nodes,num_nodes)
   A31 = zeros(Float64, num_nodes,num_nodes)
   A32 = zeros(Float64, num_nodes,num_nodes)
   A33 = zeros(Float64, num_nodes,num_nodes)

   #Sum spring terms if there are more than one on a cross-section.
   sum_kx = sum(model.kx)
   sum_kϕ = sum(model.kϕ)

   num_kx_springs = size(model.kx)[1]
   kx_ay_kx = [model.kx[i][:] .* model.ay_kx[i][:] for i=1:num_kx_springs] #multiply kx*ay_kx terms for each spring first
   sum_kx_ay_kx = sum(kx_ay_kx) #then sum them

   kx_ay_kx_ay_kx = [model.kx[i][:] .* model.ay_kx[i][:] .* model.ay_kx[i][:] for i=1:num_kx_springs]
   sum_kx_ay_kx_ay_kx = sum(kx_ay_kx_ay_kx)

   #Calculate operator quantities on LHS  AU=B.
   for i = 1:num_nodes
      A11[i,:] = model.E .* model.Iy .* Azzzz[i,:] .+ sum_kx .* AI[i,:]
      A12[i,:] = model.E .* model.Ixy .* Azzzz[i,:]
      A13[i,:] = sum_kx_ay_kx .*AI[i,:]
      A21[i,:] = model.E .* model.Ixy .* Azzzz[i,:]
      A22[i,:] = model.E .* model.Ix .* Azzzz[i,:]
      A31[i,:] = sum_kx_ay_kx .* AI[i,:]
      A33[i,:] = model.E .* model.Cw .* Azzzz[i,:] .- model.G .* model.J .* Azz[i,:] .+ sum_kx_ay_kx_ay_kx .* AI[i,:] .+ sum_kϕ .* AI[i,:] .+ model.qx .* model.ax .* AI[i,:] .- model.qy .* model.ay .* AI[i,:]
   end

   #Calculate RHS of AU=B.
   B1 = model.qx
   B2 = model.qy
   B3 = model.qx .* model.ay .+ model.qy .* model.ax

   #Reduce problem to free dof.

   #Find fixed dof.
   fixed_dof = []
   for i = 1:length(model.supports)

      if i == 1

            fixed_dof = findall(x->abs(x-model.supports[i])<=10e-6, model.z)

      elseif i != 1

            new_fixed_dof = findall(x->abs(x-model.supports[i])<=10e-6, model.z)
            fixed_dof = [fixed_dof; new_fixed_dof]

      end

   end

   #Find free dof.
   model.free_dof = setdiff(1:num_nodes,fixed_dof)

   #Partition stiffness matrix.
   model.Kff = [A11[model.free_dof,model.free_dof] A12[model.free_dof,model.free_dof] A13[model.free_dof,model.free_dof];
      A21[model.free_dof,model.free_dof] A22[model.free_dof,model.free_dof] A23[model.free_dof,model.free_dof];
      A31[model.free_dof,model.free_dof] A32[model.free_dof,model.free_dof] A33[model.free_dof,model.free_dof]]

   # model.Ksf = [A11[fixed_dof,model.free_dof] A12[fixed_dof,model.free_dof] A13[fixed_dof,model.free_dof];
   # A21[fixed_dof,model.free_dof] A22[fixed_dof,model.free_dof] A23[fixed_dof,model.free_dof];
   # A31[fixed_dof,model.free_dof] A32[fixed_dof,model.free_dof] A33[fixed_dof,model.free_dof]]

   #Define external force vector.
   model.Ff = [B1[model.free_dof]; B2[model.free_dof]; B3[model.free_dof]]

   return model

end

function residual!(R, U, K, F)

   for i=1:length(F)
    R[i] = transpose(K[i,:])*U-F[i]
   end

   return R

end

function solve(model)

   #Set up solution matrices from governing equations.
   model = governing_equations(model)

   #Define the number of nodes along the beam.
   num_nodes = length(model.z)

   #Define the deformation vectors.
   u = zeros(Float64, num_nodes)
   v = zeros(Float64, num_nodes)
   ϕ = zeros(Float64, num_nodes)

   #Define the deformation initial guess for the nonlinear solver.
   deformation_guess = model.Kff \ model.Ff

   #Solve for the beam deformations.
   solution = nlsolve((R,U) ->residual!(R, U, model.Kff, model.Ff), deformation_guess)

   #Pull the displacements and twist from the solution results.
   u[model.free_dof] = solution.zero[1:length(model.free_dof)]
   v[model.free_dof] = solution.zero[length(model.free_dof)+1:2*length(model.free_dof)]
   ϕ[model.free_dof] = solution.zero[2*length(model.free_dof)+1:3*length(model.free_dof)]

   #Add deformations to data structure.
   model.u = u
   model.v = v
   model.ϕ = ϕ

   # #Calculate reactions.
   # model.Fs = model.Ksf * solution.zero

   return model

end

end #module


text:
@testset "reified > bridge" begin
    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, b, Bin)
    @variable(m, 0 <= x[1:3] <= 5, Int)
    @constraint(m, b := {sum(x) > 10})
    optimize!(m)
    com = CS.get_inner_model(m)

    constraint = com.constraints[1]

    variables = com.search_space
    @test CS.is_constraint_solved(constraint, constraint.fct, constraint.set, [1, 3, 3, 5])
    @test !CS.is_constraint_solved(constraint, constraint.fct, constraint.set, [1, 2, 3, 5])
    @test CS.is_constraint_solved(constraint, constraint.fct, constraint.set, [0, 2, 3, 5])

    constr_indices = constraint.indices
    @test CS.still_feasible(
        com,
        constraint,
        constraint.fct,
        constraint.set,
        constr_indices[2],
        1,
    )
    @test CS.still_feasible(
        com,
        constraint,
        constraint.fct,
        constraint.set,
        constr_indices[2],
        0,
    )
    CS.fix!(com, com.search_space[constr_indices[1]], 1)
    @test !CS.still_feasible(
        com,
        constraint,
        constraint.fct,
        constraint.set,
        constr_indices[3],
        0,
    )

    # need to create a backtrack_vec to reverse pruning
    dummy_backtrack_obj = CS.BacktrackObj(com)
    dummy_backtrack_obj.step_nr = 1
    push!(com.backtrack_vec, dummy_backtrack_obj)
    # reverse previous fix
    CS.reverse_pruning!(com, 1)
    com.c_backtrack_idx = 1

    @test CS.prune_constraint!(com, constraint, constraint.fct, constraint.set)
    for ind in constr_indices[2:4]
        @test sort(CS.values(com.search_space[ind])) == [0,1,2,3,4,5]
    end
    @test sort(CS.values(com.search_space[constr_indices[1]])) == [0, 1]
    @test CS.fix!(com, com.search_space[constr_indices[1]], 1)
    # feasible but remove 0
    @test CS.prune_constraint!(com, constraint, constraint.fct, constraint.set)
    for ind in constr_indices[2:4]
        @test sort(CS.values(com.search_space[ind])) == [1, 2, 3, 4, 5]
    end
    CS.values(com.search_space[constr_indices[1]]) == [1]
end

@testset "reified > bridge is_constraint_violated" begin
    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, b, Bin)
    @variable(m, -5 <= x[1:3] <= 5, Int)
    @constraint(m, b := {sum(x) > 10})
    optimize!(m)
    com = CS.get_inner_model(m)

    constraint = com.constraints[1]

    variables = com.search_space
    @test CS.fix!(com, variables[constraint.indices[1]], 1; check_feasibility = false)
    @test CS.fix!(com, variables[constraint.indices[2]], 3; check_feasibility = false)
    @test CS.fix!(com, variables[constraint.indices[3]], 2; check_feasibility = false)
    @test CS.fix!(com, variables[constraint.indices[4]], 5; check_feasibility = false)
    @test CS.is_constraint_violated(com, constraint, constraint.fct, constraint.set)

    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, b, Bin)
    @variable(m, -5 <= x[1:3] <= 5, Int)
    @constraint(m, b := {sum(x) > 10})
    optimize!(m)
    com = CS.get_inner_model(m)

    constraint = com.constraints[1]

    variables = com.search_space
    @test CS.fix!(com, variables[constraint.indices[1]], 1; check_feasibility = false)
    @test CS.fix!(com, variables[constraint.indices[2]], 3; check_feasibility = false)
    @test CS.fix!(com, variables[constraint.indices[3]], 3; check_feasibility = false)
    @test CS.fix!(com, variables[constraint.indices[4]], 5; check_feasibility = false)
    @test !CS.is_constraint_violated(com, constraint, constraint.fct, constraint.set)

    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, b, Bin)
    @variable(m, -5 <= x[1:3] <= 5, Int)
    @constraint(m, b := {sum(x) > 10})
    optimize!(m)
    com = CS.get_inner_model(m)

    constraint = com.constraints[1]

    variables = com.search_space
    @test CS.fix!(com, variables[constraint.indices[1]], 0; check_feasibility = false)
    @test CS.fix!(com, variables[constraint.indices[2]], 3; check_feasibility = false)
    @test CS.fix!(com, variables[constraint.indices[3]], 3; check_feasibility = false)
    @test CS.fix!(com, variables[constraint.indices[4]], 5; check_feasibility = false)
    @test CS.is_constraint_violated(com, constraint, constraint.fct, constraint.set)
end

@testset "reified is_constraint_violated test" begin
    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, b, Bin)
    @variable(m, -5 <= x[1:5] <= 5, Int)
    @constraint(m, b := {x in CS.AllDifferent()})
    optimize!(m)
    com = CS.get_inner_model(m)

    constraint = com.constraints[1]

    variables = com.search_space
    @test CS.fix!(com, variables[constraint.indices[1]], 1; check_feasibility = false)
    @test CS.fix!(com, variables[constraint.indices[2]], 3; check_feasibility = false)
    @test CS.fix!(com, variables[constraint.indices[3]], 3; check_feasibility = false)
    @test CS.is_constraint_violated(com, constraint, constraint.fct, constraint.set)

    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, b, Bin)
    @variable(m, -5 <= x[1:5] <= 5, Int)
    @constraint(m, b := {x in CS.AllDifferent()})
    optimize!(m)
    com = CS.get_inner_model(m)

    constraint = com.constraints[1]

    variables = com.search_space
    @test CS.fix!(com, variables[constraint.indices[1]], 1; check_feasibility = false)
    @test CS.fix!(com, variables[constraint.indices[2]], 1; check_feasibility = false)
    @test CS.fix!(com, variables[constraint.indices[3]], 3; check_feasibility = false)
    @test !CS.is_constraint_violated(com, constraint, constraint.fct, constraint.set)

    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, b, Bin)
    @variable(m, -5 <= x[1:5] <= 5, Int)
    @constraint(m, b := {x in CS.AllDifferent()})
    optimize!(m)
    com = CS.get_inner_model(m)

    constraint = com.constraints[1]

    variables = com.search_space
    @test CS.fix!(com, variables[constraint.indices[1]], 0; check_feasibility = false)
    @test CS.fix!(com, variables[constraint.indices[2]], 3; check_feasibility = false)
    @test CS.fix!(com, variables[constraint.indices[3]], 3; check_feasibility = false)
    @test !CS.is_constraint_violated(com, constraint, constraint.fct, constraint.set)
end

@testset "reified complement prune" begin
    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, b, Bin)
    @variable(m, 0 <= x[1:3] <= 15, Int)
    @constraint(m, b := {sum(x) > 10})
    optimize!(m)
    com = CS.get_inner_model(m)

    constraint = com.constraints[1]
    complement_constraint = constraint.complement_constraint
    @test complement_constraint.set == MOI.LessThan(10.0)
    @test all(term.coefficient == 1 for term in complement_constraint.fct.terms)
    @test complement_constraint.fct.constant == 0

    variables = com.search_space
    constr_indices = constraint.indices
    # set inactive
    @test CS.fix!(com, variables[constr_indices[1]], 0; check_feasibility = false)
    # should prune complement
    @test CS.prune_constraint!(com, constraint, constraint.fct, constraint.set)

    for ind in constr_indices[2:4]
        @test sort(CS.values(com.search_space[ind])) == collect(0:10)
    end

    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, b, Bin)
    @variable(m, 0 <= x[1:3] <= 5, Int)
    @constraint(m, b := {sum(x) <= 10})
    optimize!(m)
    com = CS.get_inner_model(m)

    constraint = com.constraints[1]
    complement_constraint = constraint.complement_constraint
    @test complement_constraint.set == CPE.Strictly(MOI.LessThan(-10.0))
    @test all(term.coefficient == -1 for term in complement_constraint.fct.terms)
    @test complement_constraint.fct.constant == 0

    variables = com.search_space
    constr_indices = constraint.indices
    # set inactive
    @test CS.fix!(com, variables[constr_indices[1]], 0; check_feasibility = false)
    # should prune complement
    @test CS.prune_constraint!(com, constraint, constraint.fct, constraint.set)

    for ind in constr_indices[2:4]
        @test sort(CS.values(com.search_space[ind])) == collect(1:5)
    end

    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, b, Bin)
    @variable(m, 0 <= x[1:3] <= 5, Int)
    @constraint(m, b := {sum(x) == 10})
    optimize!(m)
    com = CS.get_inner_model(m)

    constraint = com.constraints[1]
    complement_constraint = constraint.complement_constraint
    @test complement_constraint.set == CPE.DifferentFrom(10.0)
    @test all(term.coefficient == 1 for term in complement_constraint.fct.terms)
    @test complement_constraint.fct.constant == 0

    variables = com.search_space
    constr_indices = constraint.indices
    # set inactive
    @test CS.fix!(com, variables[constr_indices[1]], 0; check_feasibility = false)
    @test CS.fix!(com, variables[constr_indices[2]], 0; check_feasibility = false)
    @test CS.fix!(com, variables[constr_indices[3]], 5; check_feasibility = false)
    # should prune complement
    @test CS.prune_constraint!(com, constraint, constraint.fct, constraint.set)

    for ind in constr_indices[4]
        @test sort(CS.values(com.search_space[ind])) == collect(0:4)
    end

    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, b, Bin)
    @variable(m, 0 <= x[1:3] <= 5, Int)
    @constraint(m, b := {sum(x) != 10})
    optimize!(m)
    com = CS.get_inner_model(m)

    constraint = com.constraints[1]
    complement_constraint = constraint.complement_constraint
    @test complement_constraint.set == MOI.EqualTo(10.0)
    @test all(term.coefficient == 1 for term in complement_constraint.fct.terms)
    @test complement_constraint.fct.constant == 0

    variables = com.search_space
    constr_indices = constraint.indices
    # set inactive
    @test CS.fix!(com, variables[constr_indices[1]], 0; check_feasibility = false)
    @test CS.fix!(com, variables[constr_indices[2]], 0; check_feasibility = false)
    @test CS.fix!(com, variables[constr_indices[3]], 5; check_feasibility = false)
    CS.changed!(com, constraint, constraint.fct, constraint.set)
    # should prune complement
    @test CS.prune_constraint!(com, constraint, constraint.fct, constraint.set)

    for ind in constr_indices[4]
        @test sort(CS.values(com.search_space[ind])) == [5]
    end
end

@testset "still_feasible reified active and inactive" begin
    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, b, Bin)
    @variable(m, 0 <= x[1:2] <= 5, Int)
    @constraint(m, b := {sum(x) > 10})
    optimize!(m)
    com = CS.get_inner_model(m)

    variables = com.search_space
    constraint = com.constraints[1]
    constr_indices = constraint.indices
    # set such that b must be 0
    @test CS.fix!(com, variables[constr_indices[2]], 0; check_feasibility = false)
    @test CS.fix!(com, variables[constr_indices[3]], 0; check_feasibility = false)

    @test !CS.still_feasible(
        com,
        constraint,
        constraint.fct,
        constraint.set,
        constr_indices[1],
        1,
    )
    @test CS.still_feasible(
        com,
        constraint,
        constraint.fct,
        constraint.set,
        constr_indices[1],
        0,
    )

    # same but b must be 1

    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, b, Bin)
    @variable(m, 0 <= x[1:2] <= 5, Int)
    @constraint(m, b := {sum(x) > 9})
    optimize!(m)
    com = CS.get_inner_model(m)

    variables = com.search_space
    constraint = com.constraints[1]
    constr_indices = constraint.indices
    # set such that b must be 0
    @test CS.fix!(com, variables[constr_indices[2]], 5; check_feasibility = false)
    @test CS.fix!(com, variables[constr_indices[3]], 5; check_feasibility = false)

    @test !CS.still_feasible(
        com,
        constraint,
        constraint.fct,
        constraint.set,
        constr_indices[1],
        0,
    )
    @test CS.still_feasible(
        com,
        constraint,
        constraint.fct,
        constraint.set,
        constr_indices[1],
        1,
    )
end

ext: jl


text:
####
# NOT FUNCTIONAL
#TODO: Convert this to v2 model - This will not work for v2
####

module v2
#Tiny Yolo V2 model configuration

using Knet: Knet, progress, progress!, gpu, KnetArray, relu, minibatch, conv4, pool, softmax

import ..sigmoid, ..Settings

mutable struct v2Chain
    layers
    v2Chain(layers...) = new(layers)
end

mutable struct Conv #Define convolutional layer
    w
    b
    stride
    padding
    f
end
mutable struct YoloPad #Define Yolo padding layer (assymetric padding).
    w
    minibatch_size::Int
end
struct Pool # Define pool layer
    size
    stride
    pad
end

YoloPad(w1::Int, w2::Int, cx::Int, cy::Int, minibatch_size::Int) =
    YoloPad(zeros(Float32, w1, w2, cx, cy), minibatch_size)#Constructor for Yolopad

Conv(w1::Int, w2::Int, cx::Int, cy::Int, st, pd, f) =
    Conv(randn(Float32, w1, w2, cx, cy), randn(Float32, 1, 1, cy, 1), st, pd, f)#Constructor for convolutional layer

#Assymetric padding function
function (y::YoloPad)(x)
    x = reshape(x, y.grid_x + 1, y.grid_y + 1, 1, 512 * y.minibatch_size)
    return reshape(conv4(y.w, x; stride = 1), y.grid_x, y.grid_y, 512, y.minibatch_size)
end

(p::Pool)(x) = pool(x; window = p.size, stride = p.stride, padding = p.pad) #pool function

(c::v2Chain)(x) = (for l in c.layers #chain function
    x = l(x)
end; x)

(c::Conv)(x) =
    c.f.(conv4(c.w, x; stride = c.stride, padding = c.padding) .+ c.b) #convolutional layer function

#leaky function
function leaky(x)
    return max(Float32(0.1) * x, x)
end

"""
    load(sets::Settings)

Loads the YOLOv2 model.

Note on Asymmetric Padding (see https://github.com/Ybakman/YoloV2#asymmetric-padding)
An alternative solution for asymmetric padding on pooling. We solved the asymmetric padding process in the following 4 steps: Let's assume we have a matrix as (d1,d2,depth,minibatch_size)

1- Apply symmetric padding = 1 with the pooling function

2- reshape it into => (d1,d2,1,minibatch_size*depth)

3- apply the following convolutional layer => 1 0  with stride = 1, padding = 0
                                              0 0

4- Lastly, reshape again into => (d1,d2,depth,minibatch_size)

Now, we made pooling with asymmetric padding on the right and bottom side. If you want to
make it on left and top side, apply this conv layer =>  0 0
                                                        0 1

The last 3 steps are implemented in YoloPad(x) function
The symmetric padding part is implemented in the pooling function just before Yolopad(x)
"""
function load(sets::Settings)
    # comments are line numbers corresponding
    # with https://github.com/pjreddie/darknet/blob/master/cfg/yolov2.cfg
return v2Chain(
    Conv(3, 3, 3, 32, 1, 1, leaky), #25
    Pool(2, 2, 0), #33
    Conv(3, 3, 32, 64, 1, 1, leaky), #37
    Pool(2, 2, 0), #45
    Conv(3, 3, 64, 128, 1, 1, leaky), #49
    Conv(1, 1, 128, 64, 1, 1, leaky), #57
    Conv(3, 3, 64, 128, 1, 1, leaky), #65
    Pool(2, 2, 0), #73
    Conv(3, 3, 128, 256, 1, 1, leaky), #77
    Conv(1, 1, 256, 128, 1, 1, leaky), #85
    Conv(3, 3, 128, 256, 1, 1, leaky), #93
    Pool(2, 2, 0), #101
    Conv(3, 3, 256, 512, 1, 1, leaky), #105
    Conv(1, 1, 512, 256, 1, 1, leaky), #113
    Conv(3, 3, 256, 512, 1, 1, leaky), #121
    Conv(1, 1, 512, 256, 1, 1, leaky), #129
    Conv(3, 3, 256, 512, 1, 1, leaky), #137
    Pool(2, 2, 0), #145
    Conv(3, 3, 512, 1024, 1, 1, leaky), #149
    Conv(1, 1, 1024, 512, 1, 1, leaky), #157
    Conv(3, 3, 512, 1024, 1, 1, leaky), #165
    Conv(3, 3, 512, 1024, 1, 1, leaky), #173
    Conv(3, 3, 512, 1024, 1, 1, leaky), #181
    ######
    Conv(3, 3, 512, 1024, 1, 1, leaky), #192
    Conv(3, 3, 512, 1024, 1, 1, leaky), #200
    #TODO: [route] #208 https://github.com/pjreddie/darknet/blob/61c9d02ec461e30d55762ec7669d6a1d3c356fb2/cfg/yolov2.cfg#L208
    Conv(1, 1, 128, 64, 1, 1, leaky), #211
    #TODO: [reorg] #219
    #TODO: [route] #222
    Conv(3, 3, 512, 1024, 1, 1, leaky), #225
    Conv(1, 1, 1024, sets.cell_bboxes * (5 + sets.num_classes), 1, 0, identity), #233
)
end

include(joinpath(@__DIR__, "loadweights.jl"))

end #module


text:
module Visualize_GeometriesWithoutMaterial

using Modia3D

filename = joinpath(Modia3D.path, "objects", "fish", "SiameseTiger0.3ds")
#filename = joinpath(Modia3D.path, "objects","engine", "crank", "crank.obj")
#file = FileShape.convexFile(filename; scaleFactor=MVector{3,Float64}(4.0,4.0,4.0))

# Properties
mat1  = Modia3D.Material(color="LightBlue", transparency=0.5)
mat2  = Modia3D.Material(transparency=0.2)
font  = Modia3D.Font(fontFamily="TimesNewRoman", bold=false, color="Black", charSize=0.2)
rtext = [0.0, 0.0, 1.0]

# Objects3D
@assembly GeometriesWithoutMaterial begin
   world            = Modia3D.Object3D(visualizeFrame=false)
   grid             = Modia3D.Object3D(world, Modia3D.Grid(1.0,0.6        ); r=[ 0.0,0.0, 0.0])
   coordinateSystem = Modia3D.Object3D(world, Modia3D.CoordinateSystem(0.7); r=[-1.5,0.0, 0.0])

   font1 = Modia3D.Font(charSize=0.2)
   font2 = Modia3D.Font(fontFamily="TimesNewRoman", bold=false, charSize=0.2, color="DarkBlue", transparency=0.0)
   text2 = Modia3D.TextShape("Text in the XY-plane"; font=font2, offset=[0.2, 0.0, 0.0],
                             axisAlignment=Modia3D.XY_Plane, alignment=Modia3D.Left)
   text3 = Modia3D.TextShape("Text in the XZ-plane"; font=font2, offset=[0.0, 0.0, 0.2],
                             axisAlignment=Modia3D.XZ_Plane, alignment=Modia3D.Left)
   text4 = Modia3D.TextShape("Text in the YZ-plane"; font=font2, offset=[0.0, 0.2, 0.0],
                             axisAlignment=Modia3D.YZ_Plane, alignment=Modia3D.Left)

   textShape2 = Modia3D.Object3D(world, text2; r=[1.2, 0.0, 0.0])
   textShape3 = Modia3D.Object3D(world, text3; r=[1.2, 0.0, 0.0])
   textShape4 = Modia3D.Object3D(world, text4; r=[1.2, 0.0, 0.0])

   # Place text above the shapes
   gridText             = Modia3D.Object3D(grid            , Modia3D.TextShape("Grid"            ; font=font); r=rtext, visualizeFrame=Modia3D.False)
   coordinateSystemText = Modia3D.Object3D(coordinateSystem, Modia3D.TextShape("CoordinateSystem"; font=font); r=rtext, visualizeFrame=Modia3D.False)
   TestUber = Modia3D.Object3D(textShape4, Modia3D.TextShape("Text"; font=font); r=rtext, visualizeFrame=Modia3D.False)

end

Modia3D.visualizeAssembly!( GeometriesWithoutMaterial(sceneOptions = Modia3D.SceneOptions(visualizeFrames=true,
                                                                                defaultFrameLength=0.7,
                                                                                enableContactDetection=false) ) )

println("... success of Visualize_GeometriesWithoutMaterial.jl!")

end


text:
export log_best, log_all

function log_best(id::String, seed::Int64, eval_count::Int64, max_fit::Float64,
                  best::Chromosome, ea::Function, ctype::DataType, log_gen::Bool)
    if log_gen
        logstr = @sprintf("E: %s %d %d %0.5f %d %d",
                          id, seed, eval_count, max_fit,
                          sum([n.active for n in best.nodes]),
                          length(best.nodes))
        if Config.log_config
            logstr = string(logstr,
                            @sprintf("%s %s %s",
                                     string(ea), string(ctype), Config.to_string()))
        end
        Logging.info(logstr)
    end
    if Config.save_best
        Logging.info(@sprintf("C: %s", string(best.genes)))
    end
    nothing
end

function log_all(id::String, seed::Int64, eval_count::Int64, max_fit::Float64,
                  best::Chromosome, ea::Function, ctype::DataType, log_gen::Bool)
    log_best(id, seed, eval_count, max_fit, best, ea, ctype, true)
end


text:
println("test MUMPS")
include("testDivGrad.jl")
include("testTwoSystem.jl")
include("testTwoSystemParallel.jl")
include("testSparseRHS.jl")
println("Done!")


text:
using UnfoldMakie
using Documenter

DocMeta.setdocmeta!(UnfoldMakie, :DocTestSetup, :(using UnfoldMakie); recursive=true)

makedocs(;
    modules=[UnfoldMakie],
    authors="Benedikt Ehinger",
    repo="https://github.com/unfoldtoolbox/UnfoldMakie.jl/blob/{commit}{path}#{line}",
    sitename="UnfoldMakie.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://unfoldtoolbox.github.io/UnfoldMakie.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
        "DesignMatrices" => "plot_design.md",
        "Results (ERP-Style)" => "plot_results.md",

    ],
)

deploydocs(;
    repo="github.com/unfoldtoolbox/UnfoldMakie.jl",
    devbranch="main",
)


text:
export pack, unpack

function pack(xs::Vector{Var}; pad=0)
    T = eltype(xs[1])
    N = ndims(xs[1])
    maxdims = zeros(Int, N)
    for x in xs
        s = size(x)
        for i = 1:N
            maxdims[i] < s[i] && (maxdims[i] = s[i])
        end
    end

    y = fill(T(pad), maxdims..., length(xs))
    st = stride(y, N+1)
    yi = 1
    for x in xs
        copyto!(y, yi, x.data, 1)
        yi += st
    end
    Var(y, (pack,xs))
end
pack(x::Node) = Node(pack, x)

function unpack(x::A, sizes::Vector) where A
    ys = map(A, sizes)
    for i = 1:length(sizes)
        ys[i] = x[s...,i]
    end
    ys
end


text:
#
# fitting.jl --
#
# Fitting Point Spread Functions (PSF) models for Julia.
#

module Fitting

import OptimPack.Powell.Newuoa
using ..PointSpreadFunctions
using ..PointSpreadFunctions:
    AbstractPSF,
    throw_bad_argument,
    to_float,
    to_int
import ..PointSpreadFunctions:
    fit,
    immutable,
    isvalid,
    parameters

# An N-dimensional Region Of Interest (ROI).
const ROI{N} = NTuple{N,Union{Colon,AbstractUnitRange{<:Integer}}}

# Union of possible model functions (may be needed in the future to avoid
# ambiguities).
const Model = Any # Union{AbstractPSF,Function}

"""
    ImmutablePSF(psf)

yields a *Point Spread Function* (PSF) model based on `psf` but whose
parameters are assumed constant.  This is useful to localize a PSF without
fitting its parameters.

"""
struct ImmutablePSF{P<:AbstractPSF} <: AbstractPSF{0}
    psf::P
end

# Just call the PSF model function when used as a function.
(P::ImmutablePSF)(args...) = P.psf(args...)

isvalid(P::ImmutablePSF) = isvalid(P.psf)
parameters(P::ImmutablePSF) = ()

ImmutablePSF(x::ImmutablePSF) = x

"""
    immutable(psf)

yields a parameter-less model of *Point Spread Function* (PSF) based on
`psf`.  This is useful to avoid fitting the parameters of a parametric
model.

"""
immutable(x::AbstractPSF) = ImmutablePSF(x)

"""
    fit(psf, pos, [wgt,] dat[, roi]; nonnegative=false) -> psf′, pos′

fits a given *Point Spread Function* (PSF) model on data `dat` with
(optional) respective weights `wgt`.  Argument `psf` is a PSF instance to
specify which PSF model to use and its initial parameters.  Argument `pos`
is the initial PSF position.  The result is a 2-tuple with the fitted PSF
model and position.

The fit is carried out by `OptimPack.Powell.Newuoa.minimize` method.  The
initial parameters should be close enough to the solution for the fit to
behave correctly.

Keyword `nonnegative` indicates whether the intensity of the PSF should be
nonnegative or not.  Keyword `rho = (rhobeg,rhoend)` may be used to specify
the initial and final precision on the parameters.  Additional keywords may
be specified and are passed to the minimizer.

An optional *Region Of Interest* (ROI) may be specified by argument `roi`
as an `N`-tuple of index ranges or colons.  This is useful to limit the
region used to perform the fit.  For example:

    fit(GaussianPSF(2.5), (23,12), dat, (10:30, :))

will fit a Gaussian PSF model on sub-array `dat[10:30,:]`.  The advantages
of specifying a ROI is that the relative position of the ROI is taken into
account in the initial and final position and that a view is used instead
of extracting a sub-array.  Specifying a ROI also works for weighted data.
Non-rectangular ROIs can be emulated by having weights equal to zero where
data should be ignored.

To simply localize the PSF without fitting its parameters, wrap the PSF
model with the `immutable` method:

    fit(immutable(psf), pos, [wgt,] dat[, roi]; nonnegative=false)

!!! note
    The `fit` and `immutable` methods are not exported (yet).  They must be
    prefixed or explicitely exported.  For example:

        using PointSpreadFunctions
        using PointSpreadFunctions: fit, immutable

"""
function fit(psf::AbstractPSF{L},
             pos::NTuple{2,Real},
             dat::AbstractArray{T,2};
             rho::Tuple{Real,Real} = (0.1, 1e-5),
             nonnegative::Bool = false,
             maxeval::Integer = 50*(L + 2),
             kwds...) where {T<:AbstractFloat,L}
    # Newuoa uses double precision floating-point.
    x = Cdouble[pos..., psf[:]...]
    ans = Newuoa.minimize!(x -> objfun(dat, psf, x, nonnegative),
                           x, rho...; maxeval=Int(maxeval), kwds...)
    return _model(psf, x), (x[1], x[2])
end

function fit(psf::AbstractPSF{L},
             pos::NTuple{2,Real},
             wgt::AbstractArray{T,2},
             dat::AbstractArray{T,2};
             rho::Tuple{Real,Real} = (0.1, 1e-5),
             nonnegative::Bool = false,
             maxeval::Integer = 50*(L + 2),
             kwds...) where {T<:AbstractFloat,L}
    @assert axes(wgt) == axes(dat)
    # Newuoa uses double precision floating-point.
    x = Cdouble[pos..., psf[:]...]
    ans = Newuoa.minimize!(x -> objfun(wgt, dat, psf, x, nonnegative),
                           x, rho...; maxeval=Int(maxeval), kwds...)
    return _model(psf, x), (x[1], x[2])
end

function fit(psf::AbstractPSF,
             pos::NTuple{N,Real},
             dat::AbstractArray{T,N},
             roi::ROI{N};
             kwds...) where {T<:AbstractFloat,N}
    off = offsets(dat, roi)
    psf1, pos1 = fit(psf, pos .- off,
                     view(dat, roi...); kwds...)
    return psf1, (pos1 .+ off)
end

function fit(psf::AbstractPSF,
             pos::NTuple{N,Real},
             wgt::AbstractArray{T,N},
             dat::AbstractArray{T,N},
             roi::ROI{N};
             kwds...) where {T<:AbstractFloat,N}
    off = offsets(dat, roi)
    psf1, pos1 = fit(psf, pos .- off,
                     view(wgt, roi...), view(dat, roi...); kwds...)
    return psf1, (pos1 .+ off)
end

offsets(A::AbstractArray{<:Any,N}, roi::ROI{N}) where {N} =
    offsets(axes(A), roi)

function offsets(inds::NTuple{N,AbstractUnitRange{<:Integer}},
                 roi::ROI{N}) where {N}
    offset(I::AbstractUnitRange{<:Integer}, ::Colon) =
        to_int(first(I)) - 1
    offset(I::AbstractUnitRange{<:Integer}, J::AbstractUnitRange{<:Integer}) =
        to_int(first(J)) - 1
    map(offset, inds, roi)
end

"""
    objfun([wgt,] dat, α, mdl, x0, y0)

yields the value of the objective function for PSF model `α*mdl(x-x0,y-y0)`
and data `dat` with optional weights `wgt`.  The objective function is
defined as:

    f(α, x0, y0) = sum_{x,y} wgt[x,y]*(dat[x,y] - α*mdl(x-x0,y-y0))^2

where the sum is evaluated for coordinates `(x,y)` taking all possible
values for the 2-dimensional array `dat`.  If specified, weights `wgt` must
all be nonnegative and have the same axes as `dat`; if not specified
`wgt[x,y] = 1` for all `(x,y)` is assumed.

The optimal value of the intensity parameter `α` may be automatically
computed given the other parameters:

    objfun([wgt,] dat, mdl, x0, y0, nonnegative=false)

yields the value of:

     min_α f(α, x0, y0) - c

where `c = f(0, x0, y0)` is an additive constant independent
of `(x0,y0)`.

As a convenience:

    objfun([wgt,] dat, trustvalues=false)

yields the value of the objective function when `α = 0`.  If optional
argument `trustvalues` is false, the data and weights values are checked
(data must have finite values and weights must be finite and nonnegative);
if `trustvalues` is true, it assumed that the data and the weights have
correct values and their values are not checked.

""" objfun

# Objective functions when alpha=0.

function objfun(dat::AbstractArray{T},
                trustvalues::Bool = false) where {T<:AbstractFloat}
    c = zero(Float64)
    if trustvalues
        @inbounds @simd for i in eachindex(dat)
            c += to_float(Float64, dat[i]^2)
        end
    else
        dflags = true
        @inbounds @simd for i in eachindex(dat)
            dflags &= isfinite(dat[i])
            c += to_float(Float64, dat[i]^2)
        end
        dflags || throw_bad_argument("there are invalid data values")
    end
    return c
end

function objfun(wgt::AbstractArray{T,N},
                dat::AbstractArray{T,N},
                trustvalues::Bool = false) where {T<:AbstractFloat,N}
    @assert axes(wgt) == axes(dat)
    c = zero(Float64)
    if trustvalues
        @inbounds @simd for i in eachindex(wgt, dat)
            c += to_float(Float64, wgt[i]*dat[i]^2)
        end
    else
        wflags = true
        dflags = true
        @inbounds @simd for i in eachindex(wgt, dat)
            wflags &= isfinite(wgt[i]) & (wgt[i] ≥ zero(wgt[i]))
            dflags &= isfinite(dat[i])
            c += to_float(Float64, wgt[i]*dat[i]^2)
        end
        wflags || throw_bad_argument("there are invalid weights")
        dflags || throw_bad_argument("there are invalid data values")
    end
    return c
end

# Objective functions when alpha is given.

function objfun(dat::AbstractArray{T,2},
                alpha::Real,
                mdl::Model,
                pos::NTuple{2,Real}) where {T<:AbstractFloat}
    objfun(dat, alpha, mdl, pos[1], pos[2])
end

function objfun(wgt::AbstractArray{T,2},
                dat::AbstractArray{T,2},
                alpha::Real,
                mdl::Model,
                pos::NTuple{2,Real}) where {T<:AbstractFloat}
    objfun(wgt, dat, alpha, mdl, pos[1], pos[2])
end

function objfun(dat::AbstractArray{T,2},
                alpha::Real,
                mdl::Model,
                x0::Real,
                y0::Real) where {T<:AbstractFloat}
    objfun(dat, to_float(T, alpha), mdl, to_float(T, x0), to_float(T, y0))
end

function objfun(wgt::AbstractArray{T,2},
                dat::AbstractArray{T,2},
                alpha::Real,
                mdl::Model,
                x0::Real,
                y0::Real) where {T<:AbstractFloat}
    objfun(wgt, dat, to_float(T, alpha), mdl, to_float(T, x0), to_float(T, y0))
end

function objfun(dat::AbstractArray{T,2},
                alpha::T,
                mdl::Model,
                x0::T,
                y0::T) where {T<:AbstractFloat}
    X, Y = axes(dat)
    c = zero(Float64)
    @inbounds for y in Y
        v = to_float(T, y) - y0
        @simd for x in X
            u = to_float(T, x) - x0
            m = mdl(u, v)
            d = dat[x,y]
            r = d - alpha*m
            c += r^2
        end
    end
    return c
end

function objfun(wgt::AbstractArray{T,2},
                dat::AbstractArray{T,2},
                alpha::T,
                mdl::Model,
                x0::T,
                y0::T) where {T<:AbstractFloat}
    X, Y = axes(dat)
    @assert axes(wgt) == (X, Y)
    c = zero(Float64)
    @inbounds for y in Y
        v = to_float(T, y) - y0
        @simd for x in X
            u = to_float(T, x) - x0
            m = mdl(u, v)
            w = wgt[x,y]
            d = dat[x,y]
            r = d - alpha*m
            c += w*r^2
        end
    end
    return c
end

# Objective functions that set alpha automatically.

function objfun(dat::AbstractArray{T,2},
                mdl::Model,
                pos::NTuple{2,Real},
                nonnegative::Bool = false) where {T<:AbstractFloat}
    objfun(dat, mdl, pos[1], pos[2], nonnegative)
end

function objfun(wgt::AbstractArray{T,2},
                dat::AbstractArray{T,2},
                mdl::Model,
                pos::NTuple{2,Real},
                nonnegative::Bool = false) where {T<:AbstractFloat}
    objfun(wgt, dat, mdl, pos[1], pos[2], nonnegative)
end

function objfun(dat::AbstractArray{T,2},
                mdl::Model,
                x0::Real,
                y0::Real,
                nonnegative::Bool = false) where {T<:AbstractFloat}
    objfun(dat, mdl, to_float(T, x0), to_float(T, y0), nonnegative)
end

function objfun(wgt::AbstractArray{T,2},
                dat::AbstractArray{T,2},
                mdl::Model,
                x0::Real,
                y0::Real,
                nonnegative::Bool = false) where {T<:AbstractFloat}
    objfun(wgt, dat, mdl, to_float(T, x0), to_float(T, y0), nonnegative)
end

function objfun(dat::AbstractArray{T,2},
                mdl::Model,
                x0::T,
                y0::T,
                nonnegative::Bool = false) where {T<:AbstractFloat}
    X, Y = axes(dat)
    a = zero(Float64)
    b = zero(Float64)
    @inbounds for y in Y
        v = to_float(T, y) - y0
        @simd for x in X
            u = to_float(T, x) - x0
            m = mdl(u, v)
            d = dat[x,y]
            a += m*m
            b += m*d
        end
    end
    if a > 0 && (nonnegative ? b > 0 : b != 0)
        # best alpha = b/a
        return -(b/a)*b
    else
        return zero(Float64)
    end
end

function objfun(wgt::AbstractArray{T,2},
                dat::AbstractArray{T,2},
                mdl::Model,
                x0::T,
                y0::T,
                nonnegative::Bool = false) where {T<:AbstractFloat}
    X, Y = axes(dat)
    @assert axes(wgt) == (X, Y)
    a = zero(Float64)
    b = zero(Float64)
    @inbounds for y in Y
        v = to_float(T, y) - y0
        @simd for x in X
            u = to_float(T, x) - x0
            m = mdl(u, v)
            w = wgt[x,y]
            d = dat[x,y]
            a += m*w*m
            b += m*w*d
        end
    end
    if a > 0 && (nonnegative ? b > 0 : b != 0)
        # best alpha = b/a
        return -(b/a)*b
    else
        return zero(Float64)
    end
end

# Objective functions designed for NEWUOA.

function objfun(dat::AbstractArray{T,2},
                psf::AbstractPSF{L},
                prm::Vector{Cdouble},
                nonnegative::Bool) where {T<:AbstractFloat,L}
    @assert length(prm) == L+2
    mdl = _model(psf, prm)
    isvalid(mdl) ?
        Cdouble(objfun(dat, mdl, prm[1], prm[2], nonnegative)) :
        zero(Cdouble)

end

function objfun(wgt::AbstractArray{T,2},
                dat::AbstractArray{T,2},
                psf::AbstractPSF{L},
                prm::Vector{Cdouble},
                nonnegative::Bool) where {T<:AbstractFloat,L}
    @assert length(prm) == L+2
    mdl = _model(psf, prm)
    isvalid(mdl) ?
        Cdouble(objfun(wgt, dat, mdl, prm[1], prm[2], nonnegative)) :
        zero(Cdouble)
end

@inline _model(psf::ImmutablePSF, prm::Vector{Cdouble}) = psf.psf

@inline _model(psf::AbstractPSF{0}, prm::Vector{Cdouble}) = psf

@inline _model(psf::T, prm::Vector{Cdouble}) where {T<:AbstractPSF{1}} =
    T(prm[3])

@inline _model(psf::T, prm::Vector{Cdouble}) where {T<:AbstractPSF{2}} =
    T(prm[3], prm[4])

end # module


text:
function addtensor!{T,N}(A::CuArray{T,N}, C::CuArray{T,N})
    h = handle(A)
    adesc = TensorDesc(A)
    cudnnAddTensor(h, T[1], adesc, A, T[1], adesc, C)
    C
end


text:
using LatBo: SandBox, Simulation
using LatBo.Playground: Feature
using LatBo.LB: Streaming, Collision, FluidKernel, LocalKernel, LocalQuantities
# the following functions are extended for mock types
import LatBo.LB.velocity
import LatBo.LB.collision
import LatBo.LB.streaming
import LatBo.LB.local_kernel

type MockCollision <: Collision
    calls
    args
    MockCollision() = new(0, nothing)
end
function collision(collider::MockCollision, args...)
    collider.calls += 1
    collider.args = deepcopy(args)
    1.5
end

type MockStreaming <: Streaming
    args :: Vector{Any}
    MockStreaming() = new(Any[])
end
function streaming(
        streamer::MockStreaming, quantities::LocalQuantities, sim::Simulation,
        from::Integer, to::Integer, direction::Integer)
    push!(streamer.args, tuple(deepcopy(quantities), sim, from, to, direction))
end

type MockKernel <: LocalKernel
    value
end

function local_kernel(kernel::MockKernel, sim::Simulation, site::Integer,
                      ::LocalQuantities, ::DenseVector)
    (sim.populations[:, site] = kernel.value)
end

facts("Local fluid kernel") do
    sim = SandBox(:D2Q9, (4, 4))

    const start = index(sim.indexing, [3, 3])
    const cᵢ = sim.lattice.celerities
    const direction = find([all(cᵢ[:, i] .== [-1, 1]) for i in 1:size(cᵢ, 2)])[1]
    const finish = index(sim.indexing, [3, 3] + cᵢ[:, direction])

    kernel = FluidKernel(
        MockCollision(),
        Dict(
            convert(Feature, 9) => MockStreaming(),
            convert(Feature, 42) => MockStreaming(),
            convert(Feature, 0) => MockStreaming()
        )
    )
    sim.playground[:] = 0
    sim.playground[finish]= 42
    sim.playground[start] = 9
    const fᵢ = 1 + rand(Float64, size(sim.populations, 1))
    const feq = equilibrium(sim.lattice, fᵢ)
    sim.populations[:, start] = fᵢ

    context("Check mock calls") do

        local_kernel(kernel, sim, start)

        @fact kernel.collision.calls --> 1
        @fact length(kernel.collision.args) --> 2
        @fact kernel.collision.args[1] --> roughly(fᵢ)
        @fact kernel.collision.args[2] --> roughly(feq)

        @fact length(kernel.streamers[42].args) --> 1
        @fact length(kernel.streamers[9].args) --> 1
        @fact length(kernel.streamers[0].args) --> size(sim.populations, 1) - 2

        @fact kernel.streamers[42].args[1][2] --> exactly(sim)
        @fact kernel.streamers[42].args[1][1].feq --> roughly(feq)
        @fact kernel.streamers[42].args[1][3:end] --> (start, finish, direction)

        @fact kernel.streamers[9].args[1][2] --> exactly(sim)
        zerodir = findfirst([all(cᵢ[:, i] .== 0) for i in 1:size(cᵢ, 2)])
        @fact kernel.streamers[9].args[1][3:end] --> (start, start, zerodir)
    end

    context("Loop over all sites") do
        sim.playground[:] = 1
        sim.playground[5] = 2
        sim.kernels = Dict(1 => MockKernel(0), 2 => MockKernel(1))

        sim.populations[:] = -1
        local_kernel(sim)

        @fact any(sim.populations == -1) --> false
        @fact sim.populations[:, 5] --> roughly(ones(sim.populations[:, 1]))
    end
end


text:
module PyBaMM

using GeneralizedGenerated
using LinearAlgebra, SparseArrays
using PyCall

include("diffeq_problems.jl")
export get_ode_problem, get_dae_problem

include("variables.jl")
export get_variable, get_l2loss_function

end # module


text:
run(`psql -f sqlsnippets/droptables.sql`)
using SearchLight
using SearchLightPostgreSQL
SearchLight.Configuration.load() |> SearchLight.connect
SearchLight.Migrations.create_migrations_table()
SearchLight.Migrations.up()


text:
#!/usr/bin/env julia

using Gasp

immutable Aelem
    idx::Int64
    rank::Int64
end

@inline nputs(s...) = ccall(:puts, Cint, (Cstring,), string("[$(grank())] ", s...))

macro tst(ex)
    oex = Expr(:inert, ex)
    quote
        r = $ex
        nputs(r ? "passed: " : "failed: ", $oex)
    end
end

if grank() == 1
    ccall(:puts, Cint, (Cstring,), string("garraytest -- $(ngranks()) ranks\n"))
end

# even distribution
# ---
nelems = ngranks() * 5

# create the array
ga = Garray(Aelem, sizeof(Aelem)+8, nelems)
@tst length(ga) == ngranks() * 5
@tst elemsize(ga) == sizeof(Aelem)+8

# get the local part
lo, hi = distribution(ga, grank())
@tst lo == ((grank()-1)*5)+1
@tst hi == lo+4

nputs(lo, "-", hi)

# write into the local part
p = access(ga, lo, hi)
nputs(hi-lo+1)
for i = 1:hi-lo+1
    p[i] = Aelem(lo+i-1, grank())
end

# let all ranks complete writing
flush(ga)
sync()

# get the whole array on rank 1 and verify it
even_dist_garray = true
if grank() == 1
    fa, fa_handle = get(ga, 1, nelems)
    for i=1:nelems
        if fa[i].idx != i
            println(i, fa[i])
            even_dist_garray = false
            break
        end
    end
    @tst even_dist_garray
end

finalize(ga)

# uneven distribution
# ---
nelems = nelems + Int(ceil(ngranks()/2))
ga = Garray(Aelem, sizeof(Aelem)+8, nelems)

# get the local part, write into it, and sync
lo, hi = distribution(ga, grank())
nputs(lo, "-", hi)
p = access(ga, lo, hi)
for i = 1:hi-lo+1
    p[i] = Aelem(lo+i-1, grank())
end
flush(ga)
sync()

# get the whole array on rank 1 and verify it
uneven_dist_garray = true
if grank() == 1
    fa, fa_handle = get(ga, 1, nelems)
    for i=1:nelems
        if fa[i].idx != i
            uneven_dist_garray = false
            break
        end
    end
    @tst uneven_dist_garray
end

finalize(ga)


text:
type CkCvode
  initStates::Array{Float64, 1}
  parameters::Array{Float64, 1}
  time::Array{Float64, 1}
  relTol::Float64
  absTol::Array{Float64, 1}

  odes::Function
  simulate_odes::Function
  simulate_data::Function

  CkCvode(odeModel::OdeModel, time, relTol, absTol) = begin
    instance = new()

    instance.initStates = collect(values(odeModel.states))
    instance.parameters = collect(values(odeModel.parameters))
    instance.time = time
    instance.relTol = relTol
    instance.absTol = absTol

    instance.odes = make_ck_cvode_odes(odeModel)
    instance.simulate_odes = (() -> simulate_ck_cvode_odes(odeModel, instance))
    instance.simulate_data = (noiseVar::Float64 -> simulate_ck_cvode_data(odeModel, instance, noiseVar))

    instance
  end
end

CkCvode(odeModel::OdeModel, time) = CkCvode(odeModel::OdeModel, time, 1e-4, [1e-8, 1e-14, 1e-6])

function codegen_ck_cvode_odes(odeModel::OdeModel)
  state_keys = collect(keys(odeModel.states))
  i = 0; state_dict = Dict{SE,SE}(Symbol(j)=>Expr(:ref, :_y, i+=1) for j in state_keys)
  i = 0; ode_dict = Dict{SE,SE}(Symbol(j)=>Expr(:ref, :_ydot, i+=1) for j in state_keys)
  i = 0; parameter_dict = Dict{SE,SE}(Symbol(j)=>Expr(:ref, :_p, i+=1) for j in keys(odeModel.parameters))

  scratch_odes = copy(odeModel.odes)
  for i in values(scratch_odes)
    replace(i, state_dict)
    replace(i, parameter_dict)
  end

  body = []
  for i in state_keys
    push!(body, :($(ode_dict[Symbol(i)]) = $(scratch_odes[i])))
  end

  @gensym ck_cvode_odes
  quote
    function $ck_cvode_odes(_t, _y, _ydot, _p)
      $(body...)
    end
  end
end

make_ck_cvode_odes(odeModel::OdeModel) = eval(codegen_ck_cvode_odes(odeModel))

function ck_cvode_ode_wrapper(t, y, ydot, user_data)
  y = Sundials.asarray(y)
  ydot = Sundials.asarray(ydot)

  user_data[1](t, y, ydot, user_data[2])

  return Int32(0)
end

function simulate_ck_cvode_odes(odeModel::OdeModel, ckCvode::CkCvode)
  nStates = length(ckCvode.initStates)
  nTimePoints = length(ckCvode.time)

  cvode_mem = Sundials.CVodeCreate(Sundials.CV_BDF, Sundials.CV_NEWTON)
  flag = Sundials.CVodeInit(cvode_mem, cfunction(ck_cvode_ode_wrapper, Int32, (Sundials.realtype, Sundials.N_Vector,
    Sundials.N_Vector, Array{Any, 1})), ckCvode.time[1], Sundials.nvector(ckCvode.initStates))
  flag = Sundials.CVodeSetUserData(cvode_mem, [ckCvode.odes, ckCvode.parameters])
  flag = Sundials.CVodeSVtolerances(cvode_mem, ckCvode.relTol, ckCvode.absTol)
  flag = Sundials.CVDense(cvode_mem, nStates)

  y = Array(Float64, nTimePoints-1, nStates)

  for i in 2:nTimePoints
    flag = Sundials.CVode(cvode_mem, (ckCvode.time)[i], ckCvode.initStates, [(ckCvode.time)[1]], Sundials.CV_NORMAL)

    if flag != Sundials.CV_SUCCESS
      throw(KeyError("SUNDIALS_ERROR: CVODE failed with flag = ", flag))
    end

    y[i-1, :] = ckCvode.initStates
  end;

  y
end

function simulate_ck_cvode_data(odeModel::OdeModel, ckCvode::CkCvode, noiseVar::Float64)
  nStates = length(ckCvode.initStates)

  simulated_odes = simulate_ck_cvode_odes(odeModel, ckCvode)
  noise = rand(MultivariateNormal(zeros(nStates), noiseVar*eye(nStates)), length(ckCvode.time)-1)'

  return simulated_odes, simulated_odes+noise
end


text:
using InfiniteOptSymbolics
using Documenter

DocMeta.setdocmeta!(InfiniteOptSymbolics, :DocTestSetup, :(using InfiniteOptSymbolics); recursive=true)

makedocs(;
    modules=[InfiniteOptSymbolics],
    authors="Joshua Pulsipher",
    repo="https://github.com/pulsipher/InfiniteOptSymbolics.jl/blob/{commit}{path}#{line}",
    sitename="InfiniteOptSymbolics.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://pulsipher.github.io/InfiniteOptSymbolics.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/pulsipher/InfiniteOptSymbolics.jl",
    devbranch="main",
)


text:
## implement Markov Chain with SEIR
include("dynamics.jl")

struct SIR <: Dynamics
    α::Float64
    β::Float64
end

# default constructor
SIR() = SIR(0.0,0.0)

# change in a day
"""
change(s::Vector{Float64},d::SIR)

Retutn the change in the state of the population in a day

s = state vector (S,I,R)
d = SIR dynamics parameters
"""
function change(s::Vector{Float64},d::SIR)
    # this is done purely for readability of the formula
    st = NamedTuple{(:S,:I,:R)}(s)
    # N = sum(s) # population size
    S = -( d.β*st.I )*st.S
    I =  ( d.β*st.I )*st.S - d.α*st.R
    R = d.α*st.R
    return [S,I,R]
end

function initialize(I::Float64,d::SIR)
    state = zeros(nstates(d))
    state[2] = I
    state[1] = 1-I
    return state
end

function nstates(d::SIR)
    return 3
end

function stateNames(d::SIR)
    return ["S" "I" "R"]
end

function getParams(α::Float64,β::Float64,d::SIR)
    return SIR(α,β)
end

"""
estimatedStates( nt::Int64, N::Int64, I0::Float64, d::SIR)

nt: time points
N: population size
I0: proportion infected at start
d: SIR parameters
"""
function estimatedStates(nt::Int64,N::Int64,I0::Float64,d::SIR)
    s0 = initialize(I0,d) # initial state
    (d,ds) = evolve(N*1.0,s0,d,nt) # evolve over time
    s1 = DataFrame(N.*d',[:S,:I,:R]) # data frame with states over time
    return s1
end

## betachange is how the β1 parameter is changed on the logit scale
function estimatedStates( N::Int64, E0::Float64,  betachange::Vector{Float64}, d::SEI3R)
    s0 = initialize(E0,d)
    (d,ds) = evolve(N*1.0,s0,betachange,d)
    s1 = DataFrame(N.*d',[:S,:I,:R])
    return s1
end


text:
struct CDARelaxation <: AbstractRelaxation
    num_layers::Int
end

function reformulate_unary_quadratic_term!(
    model::MOI.ModelLike,
    reformulation::Reformulation{CDARelaxation},
    _x::VI,
    _y::VI,
    l::Float64,
    u::Float64,
)
    num_layers = reformulation.relaxation.num_layers
    ξ_v = MOI.add_variables(model, num_layers)
    η_v = MOI.add_variables(model, num_layers)
    λ_1_v, λ_1_c = MOI.add_constrained_variables(
        model,
        [MOI.GreaterThan(0.0) for i in 1:num_layers],
    )
    λ_2_v, λ_2_c = MOI.add_constrained_variables(
        model,
        [MOI.GreaterThan(0.0) for i in 1:num_layers],
    )
    z_v, z_c = MOI.add_constrained_variables(
        model,
        [MOI.ZeroOne() for i in 1:num_layers],
    )
    x = SV(_x)
    y = SV(_y)
    ξ = [SV(ξ_v[i]) for i in 1:num_layers]
    η = [SV(η_v[i]) for i in 1:num_layers]
    λ_1 = [SV(λ_1_v[i]) for i in 1:num_layers]
    λ_2 = [SV(λ_2_v[i]) for i in 1:num_layers]
    z = [SV(z_v[i]) for i in 1:num_layers]

    θ_min = (
        if l > 0
            atan(l^2 - 1) / (2l)
        elseif l == 0
            -π / 2
        else
            atan(l^2 - 1) / (2l) - π
        end
    )
    θ_max = (
        if u > 0
            atan(u^2 - 1) / (2u)
        elseif u == 0
            -π / 2
        else
            atan(u^2 - 1) / (2u) - π
        end
    )
    @assert -3π / 2 < θ_min < π / 2
    @assert -3π / 2 < θ_max < π / 2
    θ_d = θ_max - θ_min
    θ_mid = (θ_max + θ_min) / 2

    ν = num_layers
    # RLP inequality (1)
    MOIU.normalize_and_add_constraint(
        model,
        y - (l + u) * x + l * u,
        MOI.LessThan(0.0),
    )
    # Equation (11)
    # Suggestion from Robert: Set this to equation
    MOIU.normalize_and_add_constraint(
        model,
        ξ[ν] * cos(θ_d / 2^(ν + 1)) + η[ν] * sin(θ_d / 2^(ν + 1)) -
        (y + 1.0) / 2.0 * cos(θ_d / 2^(ν + 1)),
        MOI.GreaterThan(0.0),
    )
    if !reformulation.soc_lower_bound
        # Equation 12
        MOIU.normalize_and_add_constraint(
            model,
            ξ[ν] * cos(θ_d / 2^ν) + η[ν] * sin(θ_d / 2^ν) - (y + 1.0) / 2.0,
            MOI.LessThan(0.0),
        )
        # Equation 13
        MOIU.normalize_and_add_constraint(
            model,
            ξ[ν] - (y + 1.0) / 2.0,
            MOI.LessThan(0.0),
        )
    end
    # Equation block (20)
    MOIU.normalize_and_add_constraint(
        model,
        ξ[1] - x * cos(θ_mid) - (y / 2.0 - 0.5) * sin(θ_mid),
        MOI.EqualTo(0.0),
    )
    C = max(l^2, u^2) / 2 + 0.5
    MOIU.normalize_and_add_constraint(
        model,
        (1.0λ_2[1] - 1.0λ_1[1]) * C + x * sin(θ_mid) -
        (y / 2.0 - 0.5) * cos(θ_mid),
        MOI.EqualTo(0.0),
    )
    MOI.add_constraint(
        model,
        η[1] - (1.0λ_1[1] + 1.0λ_2[1]) * C,
        MOI.EqualTo(0.0),
    )
    MOIU.normalize_and_add_constraint(
        model,
        1.0λ_1[1] - (1.0 - z[1]),
        MOI.LessThan(0.0),
    )
    MOI.add_constraint(model, 1.0λ_2[1] - 1.0z[1], MOI.LessThan(0.0))
    for j in 1:(ν-1)
        # Equation block 21
        MOI.add_constraint(
            model,
            ξ[j+1] - ξ[j] * cos(θ_d / 2^(j + 1)) - η[j] * sin(θ_d / 2^(j + 1)),
            MOI.EqualTo(0.0),
        )
        C_j = C * sin(θ_d / 2^(j + 1))
        MOI.add_constraint(
            model,
            (1.0λ_2[j+1] - 1.0λ_1[j+1]) * C_j + ξ[j] * sin(θ_d / 2^(j + 1)) -
            η[j] * cos(θ_d / 2^(j + 1)),
            MOI.EqualTo(0.0),
        )
        MOI.add_constraint(
            model,
            η[j+1] - (1.0λ_1[j+1] + 1.0λ_2[j+1]) * C_j,
            MOI.EqualTo(0.0),
        )
        MOIU.normalize_and_add_constraint(
            model,
            1.0λ_1[j] - (1.0 - z[j+1]),
            MOI.LessThan(0.0),
        )
        MOI.add_constraint(model, 1.0λ_2[j] - 1.0z[j+1], MOI.LessThan(0.0))
    end
end


text:
function ModelNet10(;mode="point_cloud")

    if mode == "point_cloud"
        return ModelNet10PCloud
    else
        error("Selected mode: $(mode) is not supported (Currently supported mode are {\"point_cloud\"}).")
    end
end

const MN10_classes_to_idx = Dict{String, UInt8}([("bathtub",1), ("bed",2), ("chair",3), ("desk",4), ("dresser",5),
    ("monitor",6), ("night_stand",7), ("sofa",8), ("table",9), ("toilet",10)])

const MN10_idx_to_classes = Dict{UInt8, String}([(1,"bathtub"), (2,"bed"), (3,"chair"), (4,"desk"), (5,"dresser"),
    (6,"monitor"), (7,"night_stand"), (8,"sofa"), (9,"table"), (10,"toilet")])

const MN10_classes = ["bathtub", "bed", "chair", "desk", "dresser",
    "monitor", "night_stand", "sofa", "table", "toilet"]

struct MN10DataPoint <: AbstractDataPoint
    idx::Int
    data::Union{PointCloud}
    ground_truth::UInt8
end

struct ModelNet10PCloud <: AbstractDataset
    root::String
    path::String #contains the path to dataset
    train::Bool
    transform #TODO Add type-assertion accordingly
    npoints::Int
    sampling #TODO Add type-assertion accordingly to include two possible option {"top", "uniform"}
    datapaths::Array
    length::Int
    classes_to_idx::Dict{String, UInt8}
    classes::Array{String,1}
end

function MN10_extract(datapath, npoints)
    cls = MN10_classes_to_idx[datapath[1]]
    pset = Array{Float32}(undef, npoints, 3)
    stream = open(datapath[2], "r")
    for i in 1:npoints
        pset[i, :] = map((x->parse(Float32, x)), split(readline(stream, keep=false), ",")[1:3])
    end
    return (pset,cls)
end

function ModelNet10PCloud(;root::String=default_root, train::Bool=true, npoints::Int=1024, transform=nothing, sampling=nothing)
    _path = normpath(dataset("ModelNet10PCloud", root))
    train ? _split="train" : _split="test"
    shapeids = [line for line in readlines(joinpath(_path, "modelnet10_$(_split).txt"))]
    shape_names = [join(split(shapeids[i], "_")[1:end-1], "_") for i in 1:length(shapeids)]
    datapaths = [(shape_names[i], joinpath(_path, shape_names[i], (shapeids[i])*".txt")) for i in 1:length(shapeids)]
    _length = length(datapaths)
    ModelNet10PCloud(root, _path, train, transform, npoints, sampling, datapaths, _length, MN10_classes_to_idx, MN10_classes)
end

function Base.getindex(v::ModelNet10PCloud, idx::Int)
    pset, cls = MN10_extract(v.datapaths[idx], v.npoints)
    return MN10DataPoint(idx, PointCloud(pset), cls)
end

Base.size(v::ModelNet10PCloud) = (v.length,)
Base.length(v::ModelNet10PCloud) = v.length


text:
export forward

using IRTools
using IRTools: var, IR, block,
    argument!, arguments, Variable,
    slots!, pis!, inlineable!,
    argnames!, varargs!, Pipe, xcall,
    insertafter!, stmt, substitute, finish,
    return!, block, arguments, blocks

"""
    Pullback{S, T}

A struct used for simulating closures. `S` is the signature of function,
`T` is a `Tuple` type to store pullbacks.
"""
struct Pullback{S, T}
    data::T
end

Pullback{S}(data::T) where {S, T} = Pullback{S, T}(data)

"""
    register(ir, F)

This function will transform the original function call in the original IR `ir` to
`forward` function calls. Besides the forward evaluation results, `forward` function
will also return a pullback function (closure) for backward evaluation use.
"""
function register(ir, F)
    pr = Pipe(ir)
    pbs = Dict{Variable, Variable}()
    argument!(pr, at = 1)
    for (v, st) in pr
        ex = st.expr
        if Meta.isexpr(ex, :call)
            yJ = insert!(pr, v, stmt(xcall(YASSAD, :forward, ex.args...), line = ir[v].line))
            pr[v] = xgetindex(yJ, 1)
            J = insertafter!(pr, v, stmt(xgetindex(yJ, 2), line = ir[v].line))
            pbs[v] = substitute(pr, J)
        end
    end
    pr = finish(pr)
    v = push!(pr, xtuple(values(pbs)...))
    pbv = push!(pr, Expr(:call, Pullback{F}, v))
    ret = pr.blocks[end].branches[end].args[1]
    ret = push!(pr, xtuple(ret, pbv))
    pr.blocks[end].branches[end].args[1] = ret
    return pr, pbs
end

"""
    forward(f, xs...)

This is the entry of IR transformation.
"""
@generated function forward(f, xs...)
    T = Tuple{f, xs...}
    m = IRTools.meta(T)
    m === nothing && return
    frw, _ = register(IR(m), T)
    argnames!(m, Symbol("#self#"), :f, :xs)
    frw = varargs!(m, frw, 2)
    # frw = slots!(pis!(inlineable!(frw)))
    return IRTools.update!(m, frw)
end

"""
    adjoint(ir, pbs)

This generates the adjoint function of forward pass. (the closure).
"""
function adjoint(ir, pbs)
    adj = empty(ir)
    self = argument!(adj)
    delta = argument!(adj)
    pullbacks = pushfirst!(adj, xcall(:getfield, self, QuoteNode(:data)))

    grads = Dict()
    grad(x, x̄) = push!(get!(grads, x, []), x̄)
    grad(x) = xaccum(adj, get(grads, x, [])...)
    grad(last(keys(ir)), delta)

    vars = keys(ir)
    for k in length(vars):-1:1
        v = vars[k]
        ex = ir[v].expr
        if haskey(pbs, v)
            pbv = insertafter!(adj, pullbacks, xcall(:getindex, pullbacks, k))
            g = push!(adj, Expr(:call, pbv, grad(v)))

            for (i, x) in enumerate(ex.args)
                x isa Variable || continue
                grad(x, push!(adj, xgetindex(g, i)))
            end
        end
    end
    gs = [grad(x) for x in arguments(ir)]
    Δ = push!(adj, xtuple(gs...))
    return!(adj, Δ)
    return adj
end

"""
our fake closure.
"""
@generated function (::Pullback{S})(delta) where S
    m = IRTools.meta(S)
    m === nothing && return
    ir = IR(m)
    _, pbs = register(ir, S)
    back = adjoint(ir, pbs)
    argnames!(m, Symbol("#self#"), :delta)
    return IRTools.update!(m, back)
end


text:
using Exodus
using Test
using Profile
using Suppressor

include("test_on_square_meshes.jl")
include("test_on_example_output.jl")
include("test_copy_and_output.jl")
# include("test_on_square_meshes_with_inclusions.jl")


text:
using LinearAlgebra, SparseArrays
using MonteCarlo
using MonteCarlo: @bm

struct State
    x::UInt16
end

const VOID = State(typemax(UInt16))
const UP = 1
const DOWN = 2

State(i::Integer) = State(UInt16(i))
function Base.getindex(s::State, i)
    mask = UInt16(1) << (i-1)
    s.x & mask == mask
end
function State(s::State, val::Bool, idx)
    mask = UInt16(1) << (idx-1)
    if val
        return State(s.x | mask)
    else
        return State(s.x & (~mask))
    end
end
function Base.show(io::IO, s::State)
    print(io, "State(", bitstring(s.x), ")")
end
function count_bits(s::State, upto=16)
    mask = VOID.x >> (16-upto)
    c = s.x & mask
    count_bits(c)
end
function count_bits(c::Integer)
    count = 0
    while c != 0
        count += 1
        c = c & (c-1)
    end
    count
end

function create(state::State, site, substate)
    lin = 2(site-1) + substate
    (state[lin] || (state == VOID)) && return 0.0, VOID

    n = count_bits(state, lin-1)
    mask = UInt16(1) << (lin-1)
    return 1.0 - 2.0 * (n%2), State(state.x | mask)
end
function annihilate(state, site, substate)
    lin = 2(site-1) + substate
    (!state[lin] || (state == VOID)) && return 0.0, VOID

    n = count_bits(state::State, lin-1)
    mask = UInt16(1) << (lin-1)
    return 1.0 - 2.0 * (n%2), State(state.x & (~mask))
end

Base.:(+)(s::State, x) = s.x + x
Base.iterate(s::State) = s == VOID ? nothing : (s, 1)
Base.iterate(s::State, _) = nothing
Base.eltype(::State) = State

function HamiltonMatrix(model::Model)
    lattice = model.l
    t = model.t
    U = - model.U
    mu = model.mu

    H = zeros(Float64, 4^length(lattice), 4^length(lattice))

    # -t ∑_ijσ c_iσ^† c_jσ
    # +U ∑_i (n_i↑ - 1/2)(n_i↓ - 1/2)
    # -μ ∑_i n_i
    for i in 1:4^length(lattice)
        lstate = State(i-1)
        for j in 1:4^length(lattice)
            rstate = State(j-1)

            E = 0.0
            # hopping (hermitian conjugate implied/included by lattice generation)
            for substate in [1, 2]
                for (source, target) in neighbors(lattice, Val(true))
                    target == -1 && continue
                    # be careful not to change rstate
                    # (or restore it with state_from_integer!)
                    _sign1, state = annihilate(rstate, source, substate)
                    _sign2, state = create(state, target, substate)
                    if _sign1 * _sign2 != 0.0 && lstate == state
                        E -= _sign1 * _sign2 * t
                    end
                end
            end

            # # U, μ terms
            for p in 1:length(lattice) #1:4
                up_occ = rstate[2(p-1)+1]
                down_occ = rstate[2(p-1)+2]
                if lstate == rstate
                    E += U * (up_occ - 0.5) * (down_occ - 0.5)
                    E -= mu * (up_occ + down_occ)
                end
            end
            H[i, j] = E
        end
    end

    # So we don't need to recalculate eigenvalues
    eigen!(H)
end

################################################################################
### equal-time Observables / Expectation values
################################################################################

# Useful two-particle operators
"""
    hopping_operator(state, site1, substate1, site2, substate2)

This is `c^\\dagger_{site2, substate2} c_{site1, substate1} |state>` i.e.
hopping from 1 to 2.
"""
@inline function _hopping_operator(state, site1, substate1, site2, substate2)
    sign1, _state = annihilate(state, site1, substate1)
    sign2, _state = create(_state, site2, substate2)
    return (sign1*sign2, _state)
end
@inline function _number_operator(state, site, substate)
    _hopping_operator(state, site, substate, site, substate)
end
@inline function _number_operator(state, site)
    p1, state1 = _number_operator(state, site, UP)
    p2, state2 = _number_operator(state, site, DOWN)
    return (p1, p2), (state1, state2)
end

function MonteCarlo.Greens(site1, site2, substate1, substate2)
    s -> begin
        _sign1, state = create(s, site1, substate1)
        # _sign1 == 0. && return 0.0, VOID
        _sign2, state = annihilate(state, site2, substate2)
        # _sign2 == 0. && return 0.0, VOID
        return _sign1 * _sign2, state
    end
end

# According to p188 QMCM this is equivalent
# can be used to check if Greens/ED is correct
function Greens_permuted(site1, site2, substate1, substate2)
    s -> begin
        delta = ((site1 == site2) && (substate1 == substate2)) ? 1.0 : 0.0
        _sign1, state = annihilate(s, site2, substate2)
        # if _sign1 == 0.0
        #     return delta, state
        # end
        if _sign1 == 0.0
            return delta == 0.0 ? (0.0, VOID) : (delta, s)
        end
        _sign2, state = create(state, site1, substate1)

        p = _sign1 * _sign2
        if p == 0.0 && delta == 0.0
            return 0.0, VOID
        else
            return delta - p, state
        end
    end
end

# Charge Density Correlation
function MonteCarlo.charge_density_correlation(site1::Integer, site2::Integer)
    state -> begin
        # states = State[]
        # prefactors = Float64[]
        # for substate2 in [UP, DOWN]
        #     sign1, _state = annihilate(state, site2, substate2)
        #     sign2, _state = create(_state, site2, substate2)
        #     sign1*sign2 == 0 && continue
        #     for substate1 in [UP, DOWN]
        #         sign3, __state = annihilate(_state, site1, substate1)
        #         sign4, __state = create(__state, site1, substate1)
        #         p = sign1 * sign2 * sign3 * sign4
        #         if p != 0.0
        #             push!(prefactors, p)
        #             push!(states, __state)
        #         end
        #     end
        # end
        # prefactors, states

        (_p1, _p2), (_state1, _state2) = _number_operator(state, site2)
        (p1, p2), (state1, state2) = _number_operator(_state1, site1)
        (p3, p4), (state3, state4) = _number_operator(_state2, site1)
        return (_p1*p1, _p1*p2, _p2*p3, _p2*p4), (state1, state2, state3, state4)
    end
end

# Magnetization
function m_x(site)
    state -> begin
        # states = typeof(state)[]
        # prefactors = Float64[]
        # _sign1, _state = annihilate(state, site, DOWN)
        # _sign2, _state = create(_state, site, UP)
        # p = _sign1 * _sign2
        # if p != 0
        #     push!(states, _state)
        #     push!(prefactors, p)
        # end
        # _sign1, _state = annihilate(state, site, UP)
        # _sign2, _state = create(_state, site, DOWN)
        # p = _sign1 * _sign2
        # if p != 0
        #     push!(states, _state)
        #     push!(prefactors, p)
        # end
        # return prefactors, states
        sign1, state1 = _hopping_operator(state, site, DOWN, site, UP)
        sign2, state2 = _hopping_operator(state, site, UP, site, DOWN)
        return (sign1, sign2), (state1, state2)
    end
end
function m_y(site)
    state -> begin
        # states = typeof(state)[]
        # prefactors = Float64[]
        # _sign1, _state = annihilate(state, site, DOWN)
        # _sign2, _state = create(_state, site, UP)
        # p = _sign1 * _sign2
        # if p != 0
        #     push!(states, _state)
        #     push!(prefactors, p)
        # end
        # _sign1, _state = annihilate(state, site, UP)
        # _sign2, _state = create(_state, site, DOWN)
        # p = _sign1 * _sign2
        # if p != 0
        #     push!(states, _state)
        #     push!(prefactors, -1.0 * p)
        # end
        # return -1im * prefactors, states
        sign1, state1 = _hopping_operator(state, site, DOWN, site, UP)
        sign2, state2 = _hopping_operator(state, site, UP, site, DOWN)
        return -1im .* (sign1, -sign2), (state1, state2)
    end
end
function m_z(site)
    state -> begin
        # states = typeof(state)[]
        # prefactors = Float64[]
        # _sign1, _state = annihilate(state, site, UP)
        # _sign2, _state = create(_state, site, UP)
        # p = _sign1 * _sign2
        # if p != 0
        #     push!(states, _state)
        #     push!(prefactors, p)
        # end
        # _sign1, _state = annihilate(state, site, DOWN)
        # _sign2, _state = create(_state, site, DOWN)
        # p = _sign1 * _sign2
        # if p != 0
        #     push!(states, _state)
        #     push!(prefactors, -1.0 * p)
        # end
        # return prefactors, states
        sign1, state1 = _number_operator(state, site, UP)
        sign2, state2 = _number_operator(state, site, DOWN)
        return (sign1, -sign2), (state1, state2)
    end
end

# Spin Density Correlations (s_{x, i} * s_{x, j} etc)
function spin_density_correlation_x(site1, site2)
    state -> begin
        # states = typeof(state)[]
        # prefactors = Float64[]
        # for substates2 in [(UP, DOWN), (DOWN, UP)]
        #     sign1, _state = annihilate(state, site2, substates2[1])
        #     sign2, _state = create(_state, site2, substates2[2])
        #     sign1 * sign2 == 0.0 && continue
        #     for substates1 in [(UP, DOWN), (DOWN, UP)]
        #         sign3, __state = annihilate(_state, site1, substates1[1])
        #         sign4, __state = create(__state, site1, substates1[2])
        #         p = sign1 * sign2 * sign3 * sign4
        #         if p != 0.0
        #             push!(prefactors, p)
        #             push!(states, __state)
        #         end
        #     end
        # end
        # prefactors, states
        _p1, _state1 = _hopping_operator(state, site2, UP, site2, DOWN)
        p1, state1 = _hopping_operator(_state1, site1, UP, site1, DOWN)
        p2, state2 = _hopping_operator(_state1, site1, DOWN, site1, UP)

        _p2, _state2 = _hopping_operator(state, site2, DOWN, site2, UP)
        p3, state3 = _hopping_operator(_state2, site1, UP, site1, DOWN)
        p4, state4 = _hopping_operator(_state2, site1, DOWN, site1, UP)

        return (_p1*p1, _p1*p2, _p2*p3, _p2*p4), (state1, state2, state3, state4)
    end
end
function spin_density_correlation_y(site1, site2)
    state -> begin
        # states = typeof(state)[]
        # prefactors = ComplexF64[]
        # for substates2 in [(UP, DOWN), (DOWN, UP)]
        #     sign1, _state = annihilate(state, site2, substates2[1])
        #     sign2, _state = create(_state, site2, substates2[2])
        #     sign1*sign2 == 0.0 && continue
        #     for substates1 in [(UP, DOWN), (DOWN, UP)]
        #         # prefactor from the - in s_y
        #         c = substates1 == substates2 ? +1.0 : -1.0
        #         sign3, __state = annihilate(_state, site1, substates1[1])
        #         sign4, __state = create(__state, site1, substates1[2])
        #         p = sign1 * sign2 * sign3 * sign4
        #         if p != 0.0
        #             push!(prefactors, -1.0 * c * p)
        #             push!(states, __state)
        #         end
        #     end
        # end
        # prefactors, states
        _p1, _state1 = _hopping_operator(state, site2, UP, site2, DOWN)
        p1, state1 = _hopping_operator(_state1, site1, UP, site1, DOWN)
        p2, state2 = _hopping_operator(_state1, site1, DOWN, site1, UP)

        _p2, _state2 = _hopping_operator(state, site2, DOWN, site2, UP)
        p3, state3 = _hopping_operator(_state2, site1, UP, site1, DOWN)
        p4, state4 = _hopping_operator(_state2, site1, DOWN, site1, UP)

        return (-_p1*p1, _p1*p2, _p2*p3, -_p2*p4), (state1, state2, state3, state4)
    end
end
function spin_density_correlation_z(site1, site2)
    state -> begin
        # states = typeof(state)[]
        # prefactors = Float64[]
        # for substates2 in [(UP, UP), (DOWN, DOWN)]
        #     sign1, _state = annihilate(state, site2, substates2[1])
        #     sign2, _state = create(_state, site2, substates2[2])
        #     sign1 * sign2 == 0.0 && continue
        #     for substates1 in [(UP, UP), (DOWN, DOWN)]
        #         # prefactor from the - in s_z
        #         c = substates1 == substates2 ? +1.0 : -1.0
        #         sign3, __state = annihilate(_state, site1, substates1[1])
        #         sign4, __state = create(__state, site1, substates1[2])
        #         p = sign1 * sign2 * sign3 * sign4
        #         if p != 0.0
        #             push!(prefactors, c * p)
        #             push!(states, _state)
        #         end
        #     end
        # end
        # prefactors, states
        _p1, _state1 = _hopping_operator(state, site2, UP, site2, UP)
        p1, state1 = _hopping_operator(_state1, site1, UP, site1, UP) # n_i↑ n_j↑
        p2, state2 = _hopping_operator(_state1, site1, DOWN, site1, DOWN) # n_i↓ n_j↑

        _p2, _state2 = _hopping_operator(state, site2, DOWN, site2, DOWN)
        p3, state3 = _hopping_operator(_state2, site1, UP, site1, UP) # n_i↑ n_j↓
        p4, state4 = _hopping_operator(_state2, site1, DOWN, site1, DOWN) # n_i↓ n_j↓

        return (_p1*p1, -_p1*p2, -_p2*p3, _p2*p4), (state1, state2, state3, state4)
    end
end

# local s-wave
# NOTE - this may no longer match the order of other observables
# Δ_i Δ_j |ψ⟩
MonteCarlo.pairing_correlation(i::Integer, j::Integer) = pairing_correlation(i, i, j, j)
# general case
# Δ_i = c_{i, ↑} c_{j, ↓} + c_{j, ↓}^† c_{i, ↑}^† (h.c.)
function MonteCarlo.pairing_correlation(i::Integer, j::Integer, k::Integer, l::Integer)
    state -> begin
        # Δ_i Δ_j = (c_{i, ↑} c_{j, ↓} + c_{j, ↓}^† c_{i, ↑}^†) *
        #           (c_{k, ↑} c_{l, ↓} + c_{l, ↓}^† c_{k, ↑}^†)

        # Δ_j = c_{k, ↑} c_{l, ↓} + c_{l, ↓}^† c_{k, ↑}^†
        sign1, _state  = annihilate(state, l, DOWN)
        sign2, _state1 = annihilate(_state, k, UP)
        p1 = sign1 * sign2
        sign1, _state  = create(state, k, UP)
        sign2, _state2 = create(_state, l, DOWN)
        p2 = sign1 * sign2

        # Δ_i = c_{i, ↑} c_{j, ↓} + c_{j, ↓}^† c_{i, ↑}^†
        sign1, _state  = annihilate(_state1, j, DOWN)
        sign2, _state11 = annihilate(_state, i, UP)
        p11 = p1 * sign1 * sign2
        sign1, _state  = create(_state1, i, UP)
        sign2, _state12 = create(_state, j, DOWN)
        p12 = p1 * sign1 * sign2

        # Δ_i = c_{i, ↑} c_{j, ↓} + c_{j, ↓}^† c_{i, ↑}^†
        sign1, _state  = annihilate(_state2, j, DOWN)
        sign2, _state21 = annihilate(_state, i, UP)
        p21 = p2 * sign1 * sign2
        sign1, _state  = create(_state2, i, UP)
        sign2, _state22 = create(_state, j, DOWN)
        p22 = p2 * sign1 * sign2

        return (p11, p12, p21, p22), (_state11, _state12, _state21, _state22)
    end
end

# = i \sum\sigma (T[trg, src] c^\dagger(trg,\sigma, \tau) c(src, \sigma, \tau) - T[src, trg] c^\dagger(src, \sigma, \tau) c(trg, \sigma \tau))
function current_density(src, trg, hopping_matrix::AbstractArray)
    state -> begin
        # states = typeof(state)[]
        # prefactors = Float64[]
        # for substate in (UP, DOWN)
        #     # T[trg, src] c^\dagger(trg,\sigma, \tau) c(src, \sigma, \tau)
        #     sign1, _state = annihilate(state, src, substate)
        #     sign2, _state = create(_state, trg, substate)
        #     if sign1*sign2 != 0.0
        #         push!(prefactors, sign1 * sign2 * hopping_matrix[trg, src])
        #         push!(states, _state)
        #     end

        #     # - T[src, trg] c^\dagger(src, \sigma, \tau) c(trg, \sigma \tau)
        #     sign1, _state = annihilate(state, trg, substate)
        #     sign2, _state = create(_state, src, substate)
        #     if sign1*sign2 != 0.0
        #         push!(prefactors, -sign1 * sign2 * hopping_matrix[src, trg])
        #         push!(states, _state)
        #     end
        # end
        # prefactors, states
        h1 = hopping_matrix[trg, src]
        h2 = -hopping_matrix[src, trg]

        p1, state1 = _hopping_operator(state, src, UP, trg, UP)
        p2, state2 = _hopping_operator(state, trg, UP, src, UP)
        p3, state3 = _hopping_operator(state, src, DOWN, trg, DOWN)
        p4, state4 = _hopping_operator(state, trg, DOWN, src, DOWN)

        return (h1*p1, h2*p2, h1*p3, h2*p4), (state1, state2, state3, state4)
    end
end

function scalarproduct(lstate::State, values::Vector, states::Vector)
    x = zero(eltype(values))
    for k in eachindex(states)
        x += (states[k] == lstate) * values[k]
        # if states[k] == lstate
        #     x += values[k]
        # end
    end
    x
end
function scalarproduct(lstate::State, values::NTuple{N, T}, states::Tuple) where {N, T}
    x = zero(T)
    for k in eachindex(states)
        x += (states[k] == lstate) * values[k]
        # if states[k] == lstate
        #     x += values[k]
        # end
    end
    x
end
scalarproduct(lstate::State, value, state::State) = (lstate == state) * value

@bm function expectation_value(
        observable::Function,
        H::Eigen;
        T=1.0, beta = 1.0 / T,
        N_sites = 4,
        N_substates = 2
    )
    vals, vecs = H
    Z = 0.0
    O = 0.0
    T = eltype(observable(State(0))[1])
    right_coefficients = zeros(T, size(vecs, 1))

    for i in eachindex(vals)
        # exp(βEᵢ)
        weight = exp(-beta * vals[i])
        Z += weight
        right_coefficients .= zero(eltype(right_coefficients))
        # |phi⟩ = |Ô|ψᵢ⟩
        for j in 1:size(vecs, 1)
            values, states = observable(State(j-1))
            for (l, state) in enumerate(states)
                state == VOID && continue
                right_coefficients[state+1] += values[l] * vecs[j, i]
            end
        end
        # ⟨ψᵢ|phi⟩
        O += weight * dot(vecs[:, i], right_coefficients)
    end

    O / Z
end

function calculate_Greens_matrix(H::Eigen, lattice; beta=1.0, N_substates=2)
    G = Matrix{Float64}(
        undef,
        length(lattice)*N_substates,
        length(lattice)*N_substates
    )
    for substate1 in 1:N_substates, substate2 in 1:N_substates
        for site1 in 1:length(lattice), site2 in 1:length(lattice)
            G[
                length(lattice) * (substate1-1) + site1,
                length(lattice) * (substate2-1) + site2
            ] = expectation_value(
                Greens(site1, site2, substate1, substate2),
                H,
                beta = beta,
                N_sites = length(lattice),
                N_substates=N_substates
            )
        end
    end
    G
end

function energy(H::Eigen; beta = 1.0)
    vals, vecs = H
    Z = 0.0
    O = 0.0

    for i in eachindex(vals)
        # E_i exp(βEᵢ)
        weight = exp(-beta * vals[i])
        Z += weight
        O += vals[i] * weight
    end

    O / Z
end

################################################################################
### unequal time
################################################################################

# ⟨a1(τ1) a2(τ2)⟩ w/ a1(τ1) = e^{τ1H} a1(τ1) e^{-τ2H}
@bm function expectation_value(
        obsτ1::Function, obsτ2::Function, H::Eigen, τ1, τ2;
        T=1.0, beta = 1.0 / T, N_sites = 4, N_substates = 2
    )
    @assert beta ≥ τ1 ≥ τ2 ≥ 0 "Time order must be respected!"
    @bm "init" begin
        O = 0.0
        # vals sorted small to large
        vals, vecs = H
        rstate = State(0)
        vals1, _ = obsτ1(rstate)
        vals2, _ = obsτ2(rstate)
        T = typeof(first(vals1) * first(vals2))
        obsτ1_mat = zeros(T, size(vecs))
        obsτ2_mat = zeros(T, size(vecs))
    end

    @bm "obs mat" begin
        @inbounds for i in eachindex(vals)
            lstate = State(i-1)
            for j in eachindex(vals)
                rstate = State(j-1)
                values, states = obsτ2(rstate)
                obsτ2_mat[i, j] = scalarproduct(lstate, values, states)

                rstate = State(j-1)
                values, states = obsτ1(rstate)
                obsτ1_mat[i, j] = scalarproduct(lstate, values, states)
            end
        end
    end

    @bm "prepare O" begin
        # A bit faster with allocations ¯\_(ツ)_/¯
        obsτ2_mat = vecs' * obsτ2_mat * vecs
        obsτ1_mat = vecs' * obsτ1_mat * vecs
        # Should be?
        # obsτ2_mat = vecs * obsτ2_mat * vecs'
        # obsτ1_mat = vecs * obsτ1_mat * vecs'

        Z = mapreduce(E -> exp(-beta * E), +, vals)
    end

    @bm "compute O" begin
        # This seems to run much faster if "prepare O" allocates ¯\_(ツ)_/¯
        # v = obsτ2_mat * exp.(-τ2 * vals)
        # v .*= exp.(-(τ1 - τ2) * vals)
        # w = obsτ1_mat * v
        # O = dot(exp.(-(beta - τ1) * vals), w)

        # Correct for τ1 = τ2
        O = 0.0
        for n in eachindex(vals), m in eachindex(vals)
            O += exp(-(beta-τ1)*vals[n]) * obsτ1_mat[n, m] *
                 exp(-(τ1-τ2)*vals[m]) * obsτ2_mat[m, n] * exp(-τ2*vals[n])
        end
    end

    # <e^{-(β - τ2) H} c_i exp(-(τ2 - τ1) H) c^\dagger_j exp(-τ1 H)>

    O / Z
end

function calculate_Greens_matrix(H::Eigen, tau1, tau2, lattice; beta=1.0, N_substates=2)
    G = Matrix{Float64}(
        undef,
        length(lattice)*N_substates,
        length(lattice)*N_substates
    )
    # Respect time order:
    swap = tau1 < tau2
    for substate1 in 1:N_substates, substate2 in 1:N_substates
        for site1 in 1:length(lattice), site2 in 1:length(lattice)
            ctau1 = s -> annihilate(s, site2, substate2)
            ctau2 = s -> create(s, site1, substate1)

            G[
                length(lattice) * (substate1-1) + site1,
                length(lattice) * (substate2-1) + site2
            ] = expectation_value(
                swap ? ctau2 : ctau1,
                swap ? ctau1 : ctau2,
                H, swap ? tau2 : tau1, swap ? tau1 : tau2,
                # s -> annihilate!(s, site2, substate2),
                # s -> create!(s, site1, substate1),
                # H, tau1, tau2,
                beta = beta,
                N_sites = length(lattice),
                N_substates = N_substates
            ) * (1 - 2swap) # minus sign from permuting c^†(τ1) c(τ2)
        end
    end
    G
end

@bm function expectation_value_integrated(
        obsτ1::Function, obsτ2::Function, H::Eigen; step = 0.1,
        T=1.0, beta = 1.0 / T, N_sites = 4, N_substates = 2
    )
    @bm "init" begin
        O = 0.0
        vals, vecs = H
        @assert eltype(vecs) <: Real
        rstate = State(0)
        vals1, _ = obsτ1(rstate)
        vals2, _ = obsτ2(rstate)
        T1 = eltype(vals1)
        T2 = eltype(vals2)
        T = (T1 <: Complex || T2 <: Complex) ? ComplexF64 : Float64
        obsτ1_mat = zeros(T, size(vecs))
        obsτ2_mat = zeros(T, size(vecs))
        X = zeros(T, size(vecs))
    end

    @bm "obs mat" begin
        @inbounds for i in eachindex(vals)
            lstate = State(i-1)
            for j in eachindex(vals)
                rstate = State(j-1)
                values, states = obsτ2(rstate)
                obsτ2_mat[i, j] = scalarproduct(lstate, values, states)

                values, states = obsτ1(rstate)
                obsτ1_mat[i, j] = scalarproduct(lstate, values, states)
            end
        end
    end

    @bm "prepare O" begin
        mul!(X, adjoint(vecs), obsτ1_mat)
        mul!(obsτ1_mat, X, vecs)

        mul!(X, obsτ2_mat, vecs)
        mul!(obsτ2_mat, transpose(X), vecs)

        X .= obsτ1_mat .* obsτ2_mat

        # obsτ2_mat = transpose(vecs) * obsτ2_mat * vecs
        # obsτ1_mat = transpose(vecs) * obsτ1_mat * vecs
        # X = obsτ1_mat .* transpose(obsτ2_mat)
        Z = mapreduce(E -> exp(-beta * E), +, vals)
    end

    @bm "compute O" begin
        for τ in beta:-step:0.5step #0:step:beta-0.5step
            o = 0.0
            for n in eachindex(vals), m in eachindex(vals)
                o += exp(-(beta-τ)*vals[n] - τ*vals[m]) * X[n, m]
            end
            O += step * o / Z
        end
    end

    O
end

number_operator(site::Integer) = state -> _number_operator(state, site)

################################################################################
### utility
################################################################################

function state2string(state)
    sub, L = size(state)
    spin = ["↑", "↓"]
    chunks = String[]
    for p in 1:L
        str = ""
        for s in 1:2
            str *= state[s, p] ? spin[s] : "⋅"
        end
        push!(chunks, str)
    end
    join(chunks, " ")
end


text:
using ElasticFDA
using Distributions
using JLD
using Base.Test

d = load("simu_data.jld");
f = d["f"];
timet = collect(linspace(0,1,size(f,1)));
f1 = f[:,1];
f2 = f[:,2];

# test smooth_data
g = smooth_data(f1,1);
@test sum(g)-sum(f1)<1e-6
smooth_data!(f1,1);
@test sum(f1)-sum(f[:,1])<1e-6

# test trapz
d = trapz(timet, timet);
@test_approx_eq(d,0.5)

# test f_to_srsf
q1 = f_to_srsf(f1, timet);
f1a = srsf_to_f(q1,timet,f1[1]);

# test optimum reparam
gam = optimum_reparam(q1,timet,q1,f1o=f1[1],f2o=f2[1]);
@test norm(gam-linspace(0,1,101)) < 1e-10

# test warping functions
qw = warp_q_gamma(timet, q1, gam);
fw = warp_f_gamma(timet, f1, gam);
@test norm(q1-qw) < 1e-10
@test norm(f1-fw) < 1e-10

# test srsf_align
out = srsf_align(f, timet);

# test align_fPCA
out1 = align_fPCA(f, timet, MaxItr=2);

# test vert_fPCA
out1 = vert_fPCA(out["qn"],timet,out["qn"]);

# test horiz_fPCA
out1 = horiz_fPCA(out["gam"], timet);

# test gauss_model
out1 = gauss_model(out["fn"], timet, out["qn"], out["gam"]);

# test pair_warping_bayes
out1 = pair_warping_bayes(f1[1:100], f2[1:100], iter=2);

# test group_warping_bayes
out1 = group_warping_bayes(f, iter=2);

# test elastic_distance
da, dp = elastic_distance(f1, f1, timet);
@test da < 1e-10
@test dp < 1e-6

# test rgam
gam = rgam(101,.1,10);

# test elastic_regression
include("test_warp_regress.jl")
timet = collect(timet);
out = elastic_regression(f, y_orig, timet, max_itr=1);
out1 = elastic_prediction(f, timet, out, y=y_orig);

# test elastic logistic regression
include("test_warp_logistic.jl")
timet = collect(timet);
out = elastic_logistic(f, y_orig, timet, max_itr=1);
out1 = elastic_prediction(f, timet, out, y=y_orig);

# test elastic m-logistic regression
include("test_warp_mlogistic.jl")
timet = collect(timet);
out = elastic_mlogistic(f, y_orig, timet, max_itr=1);
out1 = elastic_prediction(f, timet, out, y=y_orig);

d = load("beta.jld");
beta = d["beta"];

# test curve_to_q
q1 = curve_to_q(beta[:,:,1]);
beta1 = q_to_curve(q1);

# test calc_shape_dist
beta1 = beta[:,:,1];
d1 = calc_shape_dist(beta1,beta1);

# test curve_pair_align
beta2n, q2n, gam, q1 = curve_pair_align(beta1,beta1);

# test curve_geodesic
geod, geod_q = curve_geodesic(beta1,beta[:,:,2]);

# test curve functions 'O'
mu, betamean, v, q = curve_karcher_mean(beta, maxit=1);
betan, qn, betamean1, q_mu = curve_srvf_align(beta, maxit=1);
K = curve_karcher_cov(betamean, beta);
pd = curve_principal_directions(betamean, mu, K);
s = sample_shapes(mu, K);

# test curve functions 'C'
mu, betamean, v, q = curve_karcher_mean(beta, mode='C', maxit=1);
betan, qn, betamean1, q_mu = curve_srvf_align(beta, mode='C', maxit=1);
K = curve_karcher_cov(betamean, beta, mode='C');
pd = curve_principal_directions(betamean, mu, K, mode='C');
s = sample_shapes(mu, K, mode='C');

# test image functions
vars = load("image.jld");
I1 = vars["I1"];
I2 = vars["I2"];
Inew, gam = pair_align_image(I1,I2,itermax=1);