text:
# Variational Bayesian Probabilistic Matrix Factorization
#
# Reference
# ----------
#   Lim, Yew Jin, and Yee Whye Teh. "Variational Bayesian approach to movie rating prediction."
#   Proceedings of KDD cup and workshop. Vol. 7. 2007.
#
mutable struct VBPMFModel
    M::Array
    I::Int64
    J::Int64
    n::Int64
    L::Int64
    U::Array
    V::Array
    τ²::Float64
    σ²::Array
    ρ²::Array
end

function fit(model::VBPMFModel)
    M = model.M = model.M[sortperm(model.M[:, 1]), :]
    n = model.n
    I = model.I
    J = model.J
    K = size(M)[1]
    Ψ = zeros(n, n, J)
    model.σ² = ones(n)
    model.ρ² = ones(n)
    model.U = rand(I, n)
    model.V = rand(J, n)

    #1.nitialize Sj and tj for j = 1,...,J:
    S, t = initialize_S_t(model, n, J)

    for l = 1:model.L
        #2. Update Q(ui) for i = 1,...,I:
        model.U, S, t, τ²_tmp, σ²_tmp = update_U(model, Ψ, S, t, I)

        #3. Update Q(vj) for j = 1,...,J:
        model.V, Ψ, ρ²_tmp = update_V(model, Ψ, S, t, J)

        #update Learning the Variances
        model.σ², model.ρ², model.τ² = update_learning_variances(model, I, J, K, σ²_tmp, ρ²_tmp, τ²_tmp)
    end
end

function initialize_S_t(model::VBPMFModel, n, J)
    S = zeros(n, n, J)
    for j = 1:J
        S[:, :, j] = one(zeros(n, n))
    end
    t = zeros(J, n)
    return(S, t)
end

function update_U(model::VBPMFModel, Ψ, S, t, I)
    n = model.n
    M = model.M
    U = model.U
    V = model.V
    τ² = model.τ²
    σ² = model.σ²

    τ²_tmp = 0
    σ²_tmp = zeros(n)
    σ²_matrix = one(zeros(n, n)) .* (1 ./ σ²)

    for i = 1:I
        #(a) Compute Φi and ui:
        N_i = M[M[:, 1] .== i - 1, :] #user's rate item
        N_i[:, 1:2] .+= 1 #index start 1

        Φ = inv(σ²_matrix + sum((Ψ[:, :, N_i[:, 2]] .+ mean(V[N_i[:, 2], :], dims = 1)' * mean(V[N_i[:, 2], :], dims = 1)) / τ², dims = 3)[:, :, 1])
        u = Φ * sum((N_i[:, 3]' * V[N_i[:, 2], :]) / τ², dims = 1)'
        U[i, :] = u
        σ²_tmp  = σ²_tmp + diag(Φ)

        #(b) Update Sj and tj for j ∈ N(i), and discard Φi:
        idx = 1
        for j = N_i[:, 2]
            S[:, :, j] = S[:, :, j] + (Φ + U[i, :] * U[i, :]') / τ²
            t[j, :] = t[j, :] + ((N_i[idx, 3] * U[i, :]') / τ²)'
            τ²_tmp =  τ²_tmp + (N_i[idx, 3] ^ 2) - (2 * N_i[idx, 3] * (U[i, :]' * V[j, :])) + tr((Φ + U[i, :] * U[i, :]') * (Ψ[:, :, j] + V[j, :] * V[j, :]'))
            idx = idx + 1
        end
    end

    σ²_tmp = σ²_tmp + (mean(U, dims =1) .^ 2)'

    return(U, S, t, τ²_tmp, σ²_tmp)
end

function update_V(model::VBPMFModel, Ψ, S, t, J)
    V = zeros(J, model.n)
    ρ²_tmp = 0
    #3. Update Q(vj) for j = 1,...,J:
    for j = 1:J
        Ψ[:, :, j] = inv(S[:, :, j])
        V[j, :] = (Ψ[:, :, j] * t[j, :])'
        ρ²_tmp  = ρ²_tmp .+ diag(Ψ[:, :, j])
    end
    ρ²_tmp = ρ²_tmp + (mean(V, dims =1) .^ 2)'

    return(V, Ψ, ρ²_tmp)
end

function update_learning_variances(model::VBPMFModel, I, J, K, σ²_tmp, ρ²_tmp, τ²_tmp)
    σ² = (1/(I - 1)) * σ²_tmp
    ρ² = (1/(J- 1)) * ρ²_tmp
    τ² = (1/(K - 1)) * τ²_tmp
    return(σ², ρ² , τ²)
end

function predict(model::VBPMFModel, new_data)
    R_p = model.U * model.V'
    r = zeros(size(new_data)[1])
    for i in 1:size(new_data)[1]
        r[i] = R_p[new_data[i, 1] + 1, new_data[i, 2] + 1]
    end
    return r
end


text:
import InteractiveUtils # stdlib
import Pkg # stdlib
import Test # stdlib

import Random
Random.seed!(999)

@debug(string("Julia depot paths: "), Base.DEPOT_PATH)
@debug(string("Julia load paths: "), Base.LOAD_PATH)

InteractiveUtils.versioninfo(verbose=true)

@debug(string("Attempting to import PredictMD...",))
import PredictMD
@debug(string("Successfully imported PredictMD.",))
@debug(string("PredictMD version: "),PredictMD.version(),)
@debug(
    string("PredictMD package directory: "),
    PredictMD.package_directory(),
    )

@debug(string("Attempting to import PredictMDExtra...",))
import PredictMDExtra
@debug(string("Successfully imported PredictMDExtra.",))
@debug(string("PredictMDExtra version: "),PredictMDExtra.version(),)
@debug(
    string("PredictMDExtra package directory: "),
    PredictMDExtra.package_directory(),
    )

@debug(string("Julia depot paths: "), Base.DEPOT_PATH)
@debug(string("Julia load paths: "), Base.LOAD_PATH)

if group_includes_block(TEST_GROUP, TestBlockUnitTests())
    Test.@testset "Unit tests" begin
        if false
        else
            unit_test_interval_string = "[,)"
        end
        @debug(
            "unit_test_interval_string: ",
            unit_test_interval_string,
            )
        if !is_interval(unit_test_interval_string)
            throw(
                ArgumentError(
                    string(
                        "$(unit_test_interval_string) ",
                        "is not a valid interval",
                        )
                    )
                )
        end
        unit_test_interval = construct_interval(
            unit_test_interval_string
            )
        @debug(
            "unit_test_interval: ",
            unit_test_interval,
            )
        testmodulea_filename = joinpath("PredictMDTestModuleA","PredictMDTestModuleA.jl",)
        testmoduleb_filename = joinpath(
            "PredictMDTestModuleB",
            "directory1",
            "directory2",
            "directory3",
            "directory4",
            "directory5",
            "PredictMDTestModuleB.jl",
            )
        testmodulec_filename = joinpath(
            PredictMD.maketempdir(),
            "PredictMDTestModuleC.jl",
            )
        rm(testmodulec_filename; force = true, recursive = true)
        open(testmodulec_filename, "w") do io
            write(io, "module PredictMDTestModuleC end")
        end
        include(testmodulea_filename)
        include(testmoduleb_filename)
        include(testmodulec_filename)
        test_directory = dirname(@__FILE__)
        unit_test_directory = joinpath(test_directory, "unit")
        include("unit-tests-type-definitions.jl")
        for (root, dirs, files) in walkdir(unit_test_directory)
            Test.@testset "$(root)" begin
                for file in files
                    if endswith(lowercase(strip(file)), ".jl")
                        file_path = joinpath(root, file)
                        if interval_contains_x(unit_test_interval,
                                               strip(file))
                            Test.@testset "$(file_path)" begin
                                @debug("Running $(file_path)")
                                include(file_path)
                            end
                        end
                    end
                end
            end
        end
    end
end

temp_generate_examples_dir = joinpath(
    PredictMD.maketempdir(),
    "generate_examples",
    "PredictMDTEMP",
    "examples",
    )

rm(
    temp_generate_examples_dir;
    force = true,
    recursive = true,
    )

Test.@testset "Integration tests" begin
    Test.@testset "Generate examples      " begin
        Test.@test(
            temp_generate_examples_dir == PredictMD.generate_examples(
                temp_generate_examples_dir;
                scripts = true,
                include_test_statements = true,
                markdown = false,
                notebooks = false,
                execute_notebooks = false,
                )
            )
    end
    Test.@testset "Boston housing regression example (CPU)  " begin
        @debug("Testing Boston housing regression example (CPU)")
        boston_housing_tests = [
            "01_preprocess_data.jl" => TestBlockIntegration1(),
            "02_linear_regression.jl" => TestBlockIntegration1(),
            "03_random_forest_regression.jl" => TestBlockIntegration2(),
            "04_knet_mlp_regression.jl" => TestBlockIntegration2(),
            "05_compare_models.jl" => TestBlockIntegration3(),
            "06_get_model_output.jl" => TestBlockIntegration3(),
            ]
        for test_pair in boston_housing_tests
            test_file = test_pair[1]
            test_block = test_pair[2]
            if group_includes_block(TEST_GROUP, test_block)
                Test.@testset "cpu_examples/bostonhousing/src/$(test_file)" begin
                    include(
                        joinpath(
                            temp_generate_examples_dir,
                            "cpu_examples",
                            "boston_housing",
                            "src",
                            test_file,
                            )
                        )
                end
            end
        end
    end
    Test.@testset "Breast cancer biopsy classification (CPU)" begin
        @debug("Testing breast cancer biopsy classification example (CPU)")
        breast_cancer_tests = [
            "01_preprocess_data.jl" => TestBlockIntegration4(),
            "02_smote.jl" => TestBlockIntegration4(),
            "03_logistic_classifier.jl" => TestBlockIntegration4(),
            "04_random_forest_classifier.jl" => TestBlockIntegration5(),
            "05_c_svc_svm_classifier.jl" => TestBlockIntegration5(),
            "06_nu_svc_svm_classifier.jl" => TestBlockIntegration5(),
            "07_knet_mlp_classifier.jl" => TestBlockIntegration6(),
            "08_compare_models.jl" => TestBlockIntegration7(),
            "09_get_model_output.jl" => TestBlockIntegration7(),
            ]
        for test_pair in breast_cancer_tests
            test_file = test_pair[1]
            test_block = test_pair[2]
            if group_includes_block(TEST_GROUP, test_block)
                Test.@testset "cpu_examples/breastcancer/src/$(test_file)" begin
                    include(
                        joinpath(
                            temp_generate_examples_dir,
                            "cpu_examples",
                            "breast_cancer_biopsy",
                            "src",
                            test_file,
                            )
                        )
                end
            end
        end
    end
end

ext: jl


text:
### A Pluto.jl notebook ###
# v0.11.14

using Markdown
using InteractiveUtils

# ╔═╡ 2d9b943a-fd04-11ea-2072-4bffc42a2fdf
using Pkg, DrWatson

# ╔═╡ 2d9bd37a-fd04-11ea-2509-4b297c013e6d
begin
	@quickactivate "StatisticalRethinkingStan"
	using StructuralCausalModels
	using StatisticalRethinking
end

# ╔═╡ 1fbdbf76-fd03-11ea-03fb-a52d9bfd4d4c
md"## Clip-04-07-09s.jl"

# ╔═╡ fea4ccac-01d1-11eb-2456-67b005bf8512
md"### snippet 5.7"

# ╔═╡ 2d9c5a52-fd04-11ea-3ea7-c1990f74f046
DMA_1 = OrderedDict(
  :d => [:a, :m],
  :m => :a
);

# ╔═╡ 2da77b08-fd04-11ea-2ed1-adedc0ada52b
DMA_dag_1 = DAG("DMA_dag1", DMA_1)

# ╔═╡ 2dabbff6-fd04-11ea-1990-abf87ec712da
begin
	fname1 = joinpath(mktempdir(), "DMA_dag_1.dot")
	to_graphviz(DMA_dag_1, fname1)
	Sys.isapple() && run(`open -a GraphViz.app $(fname1)`)
end;

# ╔═╡ 36ca0908-01d7-11eb-2783-bb2f25bd2f17
md"### snippet 5.9"

# ╔═╡ 2db49e3c-fd04-11ea-2343-0d29f51022cc
Text(pluto_string(basis_set(DMA_dag_1)))

# ╔═╡ a8f18662-01d4-11eb-01d1-197483e27bd0
md"##### `DMA_dag_1` has no implied conditional independencies.  This means all 3 variables are associated (even after conditioning). This is a testable implications."

# ╔═╡ 2db8f448-fd04-11ea-0ee9-1ff203ccb3cb
adjustment_sets(DMA_dag_1, :a, :d)

# ╔═╡ cbdb6fea-01d5-11eb-0ce3-131585c3339b
md"##### `DMA_dag_1` has no adjustment sets when regressing `:d` on `:a`."

# ╔═╡ 4458c102-01d7-11eb-3d8c-93f9129401fc
md"### snippet 5.8"

# ╔═╡ 2dbfc582-fd04-11ea-2086-e1f761ff9365
DMA_2 = OrderedDict(
  [:d, :m] => :a
)

# ╔═╡ 2dc0543e-fd04-11ea-39c9-6d3da3587f12
DMA_dag_2 = DAG("DMA_dag_2", DMA_2)

# ╔═╡ 2dcccbba-fd04-11ea-104d-797e2b94ebbb
begin
	fname2 = joinpath(mktempdir(), "DMA_dag_2.dot")
	to_graphviz(DMA_dag_2, fname2)
	Sys.isapple() && run(`open -a GraphViz.app $(fname2)`)
end;

# ╔═╡ 18f56d12-fd10-11ea-02a2-c77de8c1416b
Text(pluto_string(basis_set(DMA_dag_2)))

# ╔═╡ f199e5a2-01d5-11eb-3b15-95701444ac4d
md"##### `DMA_dag_2` has 1 testable conditional independence (:d and :m are independent after conditioning on :a). Otherwise all pairs of variables are associated"

# ╔═╡ 2dd8856a-fd04-11ea-1486-2f3c7dd650bb
adjustment_sets(DMA_dag_2, :a, :d)

# ╔═╡ 104634ec-01d6-11eb-186a-7f870a0f4cb8
md"##### `DMA_dag_2` has no adjustment sets."

# ╔═╡ 2de11e4e-fd04-11ea-3d2a-3b1a3dbaeb63
md"## End of clip-05-07-09s.jl"

# ╔═╡ Cell order:
# ╟─1fbdbf76-fd03-11ea-03fb-a52d9bfd4d4c
# ╠═2d9b943a-fd04-11ea-2072-4bffc42a2fdf
# ╠═2d9bd37a-fd04-11ea-2509-4b297c013e6d
# ╟─fea4ccac-01d1-11eb-2456-67b005bf8512
# ╠═2d9c5a52-fd04-11ea-3ea7-c1990f74f046
# ╠═2da77b08-fd04-11ea-2ed1-adedc0ada52b
# ╠═2dabbff6-fd04-11ea-1990-abf87ec712da
# ╟─36ca0908-01d7-11eb-2783-bb2f25bd2f17
# ╠═2db49e3c-fd04-11ea-2343-0d29f51022cc
# ╟─a8f18662-01d4-11eb-01d1-197483e27bd0
# ╠═2db8f448-fd04-11ea-0ee9-1ff203ccb3cb
# ╠═cbdb6fea-01d5-11eb-0ce3-131585c3339b
# ╟─4458c102-01d7-11eb-3d8c-93f9129401fc
# ╠═2dbfc582-fd04-11ea-2086-e1f761ff9365
# ╠═2dc0543e-fd04-11ea-39c9-6d3da3587f12
# ╠═2dcccbba-fd04-11ea-104d-797e2b94ebbb
# ╠═18f56d12-fd10-11ea-02a2-c77de8c1416b
# ╠═f199e5a2-01d5-11eb-3b15-95701444ac4d
# ╠═2dd8856a-fd04-11ea-1486-2f3c7dd650bb
# ╟─104634ec-01d6-11eb-186a-7f870a0f4cb8
# ╟─2de11e4e-fd04-11ea-3d2a-3b1a3dbaeb63


text:
using Pkg
Pkg.activate((@__DIR__)*"/..")

using Plasmo
using MPI
using PipsNLP

MPI.Init()

graph = OptiGraph()

@optinode(graph,n1)
@variable(n1,x0 >= 0)

subgraph = OptiGraph()
add_subgraph!(graph,subgraph)
@optinode(subgraph,n2)
@optinode(subgraph,n3)
@optinode(subgraph,n4)

#Linked variables using linkconstraints
@variable(n2,x >= 0)
@variable(n3,x >= 0)
@variable(n4,x >= 0)

#Local models
if rank == 0
    @variable(n2,x2 >= 2)
    @variable(n2,y2 >= 0)
    @constraint(n2,x2 + y2 == 3)
    @NLconstraint(n2,x2^3 <= 10)
    @constraint(n2,n2[:x] + n2[:x2] <= 10)
    @linkconstraint(graph,n1[:x0] == y2)
    @objective(n2,Min,x2)
elseif rank == 1
    @variable(n3,x3 >= 2)
    @variable(n3,y3 >= 0)
    @constraint(n3,x3 + y3 == 4)
    @constraint(n3,n3[:x] + n3[:x3] <= 10)
    @linkconstraint(graph,n1[:x0] == y3)
    @objective(n3,Min,x3)
elseif rank == 2
    @variable(n4,x4 >= 2)
    @variable(n4,y4 >= 0)
    @constraint(n4,x4 + y4 == 5)
    @constraint(n4,n4[:x] + n4[:x4] <= 10)
    @linkconstraint(graph,n1[:x0] == y4)
    @objective(n4,Min,x4)
end
#Link constraints between subnodes
@linkconstraint(graph,n2[:x] == n3[:x])
@linkconstraint(graph,n3[:x] == n4[:x])

#Solve with PIPS-NLP
PipsNLP.pipsnlp_solve(graph)

@mpi_do manager begin
    if rank == 0
        println("first stage: ")
        @show (value(n1[:x0]))

        println("second stage: n2")
        @show (value(n2[:y2]))
        @show (value(n2[:x]))
    end
end

@mpi_do manager begin
    if rank == 1
        println("second stage: n3")
        @show (value(n3[:y3]))
        @show (value(n3[:x]))
    end
end

@mpi_do manager begin
    if rank == 2
        println("second stage: n4")
        @show (value(n4[:y4]))
        @show (value(n4[:x]))
    end
end

MPI.finalize()


text:
#############################################################################
# cutting_plane.jl
# Implement the cutting plane method to solve convex optimization of the form:
# minimize f(x) + g(x), where f(x) is the Lovasz extension of a submodular funciton.
#############################################################################

# using SCS
# using ECOS
#using Gurobi
using Mosek

export cutting_plane

function cutting_plane(p::Problem,
                       f::LovaszExtAtom;
                       s::AbstractMathProgSolver = MosekSolver(),
                       abs_tol::Float64 = 1e-3,
                       max_iters::Int = 100,
                       max_rep::Int = 100,
                       lev_tol_def::Float64 = 1e-3,
                       lev_tol_up::Float64 = 1e-1,
                       lev_tol_low::Float64 = 1e-7,
                       shrink_point::Int = 20,
                       λ::Float64 = 1.5,
                       μ::Float64 = 0.7)

  lev_tol = lev_tol_def

  if length(get_v(p.objective + f)) != 1
    error("Cannot optimize functions with multiple variables.")
  end
  q = Problem(p.head, p.objective, p.constraints)
  conslength = length(p.constraints)

  # solve!(q, GurobiSolver(OutputFlag=0, OptimalityTol = 1e-3))
  # solve!(q, SCSSolver(verbose=false))
  # solve!(q, ECOSSolver(verbose=false, abstol = 1e-6))
  solve!(q, MosekSolver(MSK_IPAR_LOG = 0, MSK_DPAR_INTPNT_CO_TOL_REL_GAP = 1e-3 * abs_tol))

  # Reset the primal variable for initialization
  push!(q.solution.primal, q.solution.primal[end])
  q.solution.primal[end - 1] = 0

  var = get_v(q.objective)[1]
  t = Variable(1)
  q.objective += t
  n = size(var)[1]
  a = greedy(f.func, var.value[:])
  push!(q.constraints, dot(a, var) <= t)
  constraints = []
  push!(constraints, q.constraints[end].lhs)       # the array of the affine functions
  subgradients = [a]

  cur_upper = 0.0
  cur_low = 0.0
  upper = Inf
  lower = -Inf
  optsol = var.value[:]
  pernum = 1
  permset = [sortperm(optsol, rev = true)]         # orderings explored
  low_sol = zeros(length(var.value[:]))
  iters = 1
  rep = 0                                          # the repetition count of a trial point
  cos_num = 0.0
  gap = 0.0
  gapcount = []

  qq = Problem(q.head, q.objective, q.constraints)

  for i = 1:max_iters

    # println(length(permset))

    # solve!(q, GurobiSolver(OutputFlag=0, OptimalityTol = TOL^2), warmstart = true)
    # solve!(q, SCSSolver(verbose=false, suppress_warnings = true), warmstart = true)
    # solve!(q, ECOSSolver(verbose=false, abstol = 1e-6))
    solve!(q, MosekSolver(MSK_IPAR_LOG = 0, MSK_DPAR_INTPNT_CO_TOL_REL_GAP = 1e-3 * abs_tol))

    if isnan(q.optval)
      var.value[:] = optsol
      # println("Mosek returns NaN!")
      # println("gapopt = $gap")
      break
    end

    inacct = t.value                             # the value of t returned by the solver
    acct = maximum(evaluate.(constraints))       # the value of f_{(i)}(x^{(i)})
    cur_sol = var.value[:]
    cur_upper = q.optval - inacct + evaluate(f)[1]
    if cur_upper < upper
      optsol = var.value[:]
      upper = copy(cur_upper)
    end

    if q.status == :Optimal
      cur_low = q.optval - inacct + acct
    else
      cur_low = solvedualMosek!(q)
    end

    if lower < cur_low                           # update the lower bound
      low_sol = var.value[:]
    # else
      # println("lowstab")
    end
    lower = max(cur_low, lower)
    gap = upper - lower
    push!(gapcount, gap)
    if gap < abs_tol
      var.value[:] = optsol
      # println("total = $i")
      # println("gapopt = $gap")
      break
    else
      # println("i = $i")
      # println("gapsub = $gap")
    end
    if upper < cur_upper
      rep += 1
      if length(permset) >= max_rep
        var.value[:] = optsol
        println("reached the maximum iteration at a single point")
        # println("gapsub = $gap")
        break
      end
      if rep == shrink_point                         # check if the solver is stuck
        while pernum > max_rep && lev_tol >= lev_tol_low  # println("reduce permnum!")
          lev_tol *= μ
          (c, pernum, perm) = greedy_rand(f.func, optsol, lev_tol)
        end
      end
      b = greedy(f.func, var.value[:])
      while in(b, subgradients)
        # println("b included!")
        (b, pernum1, perm) = greedy_rand(f.func, var.value[:], lev_tol)
      end
      var.value[:] = optsol
      while length(permset) >= pernum && lev_tol <= lev_tol_up  # println("enlarge permnum!")
        lev_tol *= λ
        (a, pernum, perm) = greedy_rand(f.func, optsol, lev_tol)
      end
      (a, pernum, perm) = greedy_rand(f.func, optsol, lev_tol)
      while in(a, subgradients)
        # println("a included!")
        if !in(perm, permset)
          push!(permset, perm)
        end
        (a, pernum) = greedy_rand(f.func, optsol, lev_tol)
      end
      push!(permset, perm)
    else                                           # update the upper bound
      rep = 1                                      # reset repetition count
      lev_tol = lev_tol_def                        # reset the level set tolerance
      (a, pernum, perm) = greedy_rand(f.func, optsol, lev_tol)
      a = greedy(f.func, optsol)
      permset = [sortperm(var.value[:], rev = true)]
    end

    cos_num = length(constraints)
    if cos_num > n
      prod = zeros(cos_num)
      for i = 1:cos_num
        prod[i] = dot(optsol, subgradients[i])
      end
      if cos_num == (2 * n + 3)
        minn = indmin(prod)
        minn = indmin(prod)
        q.constraints[conslength + minn] = q.constraints[end]
        constraints[minn] = constraints[end]
        subgradients[minn] = subgradients[end]
        prod = prod[1:end-1]
        q.constraints = q.constraints[1:end-1]
        constraints = constraints[1:end-1]
        subgradients = subgradients[1:end-1]
      end
      if rep == 1
        minn = indmin(prod)
        q.constraints[conslength + minn] = (dot(a, var) <= t)
        constraints[minn] = q.constraints[conslength + minn].lhs
        subgradients[minn] = a
      else                                       # b is added
        minn = indmin(prod)
        q.constraints[conslength + minn] = (dot(a, var) <= t)
        constraints[minn] = q.constraints[conslength + minn].lhs
        subgradients[minn] = a
        maxx = indmin(prod)
        prod[minn] = prod[maxx]
        minn1 = indmin(prod)
        q.constraints[conslength + minn1] = (dot(b, var) <= t)
        constraints[minn1] = q.constraints[conslength + minn1].lhs
        subgradients[minn1] = b
      end
    else
      if rep > 1
        push!(q.constraints, dot(b, var) <= t)
        push!(constraints, q.constraints[end].lhs)
        push!(subgradients, b)
      end
      push!(q.constraints, dot(a, var) <= t)
      push!(constraints, q.constraints[end].lhs)
      push!(subgradients, a)
    end
    iters += 1
  end
  # return optsol, iters, gapcount
  return gapcount
end

function cutting_plane(p::Problem,
                       f::LovaszExtAbsAtom;
                       s::AbstractMathProgSolver = MosekSolver(),
                       abs_tol::Float64 = 1e-3,
                       max_iters::Int = 100,
                       max_rep::Int = 100,
                       lev_tol_def::Float64 = 1e-3,
                       lev_tol_up::Float64 = 1e-1,
                       lev_tol_low::Float64 = 1e-7,
                       shrink_point::Int = 20,
                       λ::Float64 = 1.5,
                       μ::Float64 = 0.7)

  lev_tol = lev_tol_def

  if length(get_v(p.objective + f)) != 1
    error("Cannot optimize functions with multiple variables.")
  end
  q = Problem(p.head, p.objective, p.constraints)
  conslength = length(p.constraints)

  # solve!(q, GurobiSolver(OutputFlag=0, OptimalityTol = 1e-3))
  # solve!(q, SCSSolver(verbose=false))
  # solve!(q, ECOSSolver(verbose=false, abstol = 1e-6))
  solve!(q, MosekSolver(MSK_IPAR_LOG = 0, MSK_DPAR_INTPNT_CO_TOL_REL_GAP = 1e-3 * abs_tol))

  # Reset the primal variable for warmstart
  push!(q.solution.primal, q.solution.primal[end])
  q.solution.primal[end - 1] = 0

  var = get_v(q.objective)[1]
  t = Variable(1)
  q.objective += t
  n = size(var)[1]
  a = greedy(f.func, abs.(var.value[:]))
  push!(q.constraints, dot(a, abs(var)) <= t)
  constraints = []
  push!(constraints, q.constraints[end].lhs)       # the array of the affine functions
  subgradients = [a]

  cur_upper = 0.0
  cur_low = 0.0
  upper = Inf
  lower = -Inf
  optsol = var.value[:]
  pernum = 1
  permset = [sortperm(abs.(optsol), rev = true)]         # orderings explored
  low_sol = zeros(length(var.value[:]))
  iters = 1
  rep = 0                                          # the repetition count of a trial point
  cos_num = 0.0
  gap = 0.0

  for i = 1:max_iters

    # println(length(permset))

    # solve!(q, GurobiSolver(OutputFlag=0, OptimalityTol = TOL^2), warmstart = true)
    # solve!(q, SCSSolver(verbose=false, suppress_warnings = true), warmstart = true)
    # solve!(q, ECOSSolver(verbose=false, abstol = 1e-6))
    solve!(q, MosekSolver(MSK_IPAR_LOG = 0, MSK_DPAR_INTPNT_CO_TOL_REL_GAP = 1e-3 * abs_tol))

    if isnan(q.optval)
      var.value[:] = optsol
      # println("Mosek returns NaN!")
      # println("gapopt = $gap")
      break
    end

    inacct = t.value                             # the inaccurate value of t
    acct = maximum(evaluate.(constraints))       # the accurate value of t
    cur_sol = var.value[:]
    cur_upper = q.optval - inacct + evaluate(f)[1]
    if cur_upper < upper
      optsol = var.value[:]
      upper = copy(cur_upper)
    end

    if q.status == :Optimal
      cur_low = q.optval - inacct + acct
    else
      cur_low = solvedualMosek!(q)
    end

    if lower < cur_low                           # update the lower bound
      low_sol = var.value[:]
    # else
      # println("lowstab")
    end
    lower = max(cur_low, lower)
    gap = upper - lower
    if gap < abs_tol
      var.value[:] = optsol
      # println("gapopt = $gap")
      break
    # else
      # println("gapsub = $gap")
    end
    if upper < cur_upper
      rep += 1
      if length(permset) >= max_rep
        var.value[:] = optsol
        println("reached the maximum iteration at a single point")  # println("gapsub = $gap")
        break
      end
      if rep == shrink_point                         # check if the solver is stuck
        while pernum > max_rep && lev_tol >= lev_tol_low  # println("reduce permnum!")
          lev_tol *= μ
          (c, pernum, perm) = greedy_rand(f.func, abs.(optsol), lev_tol)
        end
      end
      b = greedy(f.func, abs.(var.value[:]))
      while in(b, subgradients)
        # println("b included!")
        (b, pernum1, perm) = greedy_rand(f.func, abs.(var.value[:]), lev_tol)
      end
      var.value[:] = optsol
      while length(permset) >= pernum && lev_tol <= lev_tol_up  # println("enlarge permnum!")
        lev_tol *= λ
        (a, pernum, perm) = greedy_rand(f.func, abs.(optsol), lev_tol)
      end
      (a, pernum, perm) = greedy_rand(f.func, abs.(optsol), lev_tol)
      while in(a, subgradients)
        # println("a included!")
        if !in(perm, permset)
          push!(permset, perm)
        end
        (a, pernum) = greedy_rand(f.func, abs.(optsol), lev_tol)
      end
      push!(permset, perm)
    else                                           # update the upper bound
      rep = 1                                      # reset repetition count
      lev_tol = lev_tol_def                        # reset the level set tolerance
      (a, pernum, perm) = greedy_rand(f.func, abs.(optsol), lev_tol)
      a = greedy(f.func, abs.(optsol))
      permset = [sortperm(var.value[:], rev = true)]
    end

    cos_num = length(constraints)
    if cos_num > n
      prod = zeros(cos_num)
      for i = 1:cos_num
        prod[i] = dot(optsol, subgradients[i])
      end
      if cos_num == (2 * n + 3)
        minn = indmin(prod)
        minn = indmin(prod)
        q.constraints[conslength + minn] = q.constraints[end]
        constraints[minn] = constraints[end]
        subgradients[minn] = subgradients[end]
        prod = prod[1:end-1]
        q.constraints = q.constraints[1:end-1]
        constraints = constraints[1:end-1]
        subgradients = subgradients[1:end-1]
      end
      if rep == 1
        minn = indmin(prod)
        q.constraints[conslength + minn] = (dot(a, abs(var)) <= t)
        constraints[minn] = q.constraints[conslength + minn].lhs
        subgradients[minn] = a
      else                                       # b is added
        minn = indmin(prod)
        q.constraints[conslength + minn] = (dot(a, abs(var)) <= t)
        constraints[minn] = q.constraints[conslength + minn].lhs
        subgradients[minn] = a
        maxx = indmin(prod)
        prod[minn] = prod[maxx]
        minn1 = indmin(prod)
        q.constraints[conslength + minn1] = (dot(b, abs(var)) <= t)
        constraints[minn1] = q.constraints[conslength + minn1].lhs
        subgradients[minn1] = b
      end
    else
      if rep > 1
        push!(q.constraints, dot(b, abs(var)) <= t)
        push!(constraints, q.constraints[end].lhs)
        push!(subgradients, b)
      end
      push!(q.constraints, dot(a, abs(var)) <= t)
      push!(constraints, q.constraints[end].lhs)
      push!(subgradients, a)
    end
    iters += 1
  end
  return optsol
end


text:
"""
```
removeNode!(network::DislocationNetwork, nodeGone, lastNode = nothing)
```
Removes node `nodeGone` from network.
"""
function removeNode!(network::DislocationNetwork, nodeGone, lastNode = nothing)
    links = network.links
    coord = network.coord
    label = network.label
    nodeVel = network.nodeVel
    nodeForce = network.nodeForce
    connectivity = network.connectivity

    isnothing(lastNode) ? lastNode = maximum((network.numNode[1], 1)) : nothing

    # The if nodeGone is not the last node in the relevant arrays, replace it by lastNode.
    if nodeGone < lastNode
        label[nodeGone] = label[lastNode]
        @inbounds @simd for i in 1:3
            coord[i, nodeGone] = coord[i, lastNode]
            nodeVel[i, nodeGone] = nodeVel[i, lastNode]
            nodeForce[i, nodeGone] = nodeForce[i, lastNode]
        end
        @inbounds @simd for i in 1:size(connectivity, 1)
            connectivity[i, nodeGone] = connectivity[i, lastNode]
        end
        # Change the link
        @inbounds @simd for i in 1:connectivity[1, nodeGone]
            idx = 2 * i
            links[connectivity[idx + 1, nodeGone], connectivity[idx, nodeGone]] = nodeGone
        end
    end

    # Remove the last node from the arrays since nodeGone was either replaced by lastNode above, or already was the last node. Reduce the number of nodes by one.
    coord[:, lastNode] .= 0
    label[lastNode] = 0
    nodeVel[:, lastNode] .= 0
    nodeForce[:, lastNode] .= 0
    network.numNode[1] -= 1
    connectivity[:, lastNode] .= 0

    return network
end

"""
```
removeConnection!(network::DislocationNetwork, nodeKept, connectGone)
```
Removes connection `connectGone` from `nodeKept`.
"""
function removeConnection!(network::DislocationNetwork, nodeKept, connectGone)
    connectivity = network.connectivity
    linksConnect = network.linksConnect

    # Remove entry from linksConnect.
    idx = 2 * connectGone
    link1 = connectivity[idx, nodeKept]
    @assert link1 != 0 "removeConnection!: connection $connectGone for node $nodeKept does not exist."
    link2 = connectivity[idx + 1, nodeKept]
    linksConnect[link2, link1] = 0

    # The last connection to be added to nodeKept is equal to the connectivity of the nodeKept.
    lastConnect = connectivity[1, nodeKept]
    lst = 2 * lastConnect

    # If the connectGone is not the last connection that was made, replace it by the last connection.
    if connectGone < lastConnect
        connectivity[idx, nodeKept] = connectivity[lst, nodeKept]
        connectivity[idx + 1, nodeKept] = connectivity[lst + 1, nodeKept]
        linksConnect[connectivity[idx + 1, nodeKept], connectivity[idx, nodeKept]] =
            connectGone
    end

    # Change connectivity to reflect that nodeKept has one less connection.
    # Remove the last connection since lastConnect was either replaced by connectGone above, or already was the last connection.
    connectivity[1, nodeKept] -= 1
    connectivity[lst:(lst + 1), nodeKept] .= 0

    return network
end

"""
```
removeLink!(network::DislocationNetwork, linkGone, lastLink = nothing)
```
Removes link `linkGone` and uses `lastLink` to reoganise the network.
"""
function removeLink!(network::DislocationNetwork, linkGone, lastLink = nothing)
    links = network.links
    slipPlane = network.slipPlane
    bVec = network.bVec
    segForce = network.segForce
    connectivity = network.connectivity
    linksConnect = network.linksConnect

    @assert linkGone <= size(links, 2) "removeLink!: link $linkGone not found."

    # Delete linkGone from connectivity for both nodes in the link.
    node1 = links[1, linkGone]
    connectGone1 = linksConnect[1, linkGone]
    removeConnection!(network, node1, connectGone1)

    node2 = links[2, linkGone]
    connectGone2 = linksConnect[2, linkGone]
    removeConnection!(network, node2, connectGone2)

    # Remove link that no longer appears in connectivity and doesn't connect any nodes.
    checkLink = SVector{2, Int}(linksConnect[1, linkGone], linksConnect[2, linkGone])
    @assert checkLink ⋅ checkLink == 0 "removeLink!: link $linkGone still has connections and should not be deleted."

    isnothing(lastLink) ? lastLink = maximum((network.numSeg[1], 1)) : nothing

    # If the linkGone is not the last link, replace it with lastLink.
    if linkGone < lastLink
        @inbounds @simd for i in 1:2
            links[i, linkGone] = links[i, lastLink]
            linksConnect[i, linkGone] = linksConnect[i, lastLink]
        end
        @inbounds @simd for i in 1:3
            slipPlane[i, linkGone] = slipPlane[i, lastLink]
            bVec[i, linkGone] = bVec[i, lastLink]
        end
        @inbounds for j in 1:2
            @simd for i in 1:3
                segForce[i, j, linkGone] = segForce[i, j, lastLink]
            end
        end
        node1 = links[1, linkGone]
        node2 = links[2, linkGone]
        connect1 = 2 * linksConnect[1, linkGone]
        connect2 = 2 * linksConnect[2, linkGone]
        connectivity[connect1, node1] = linkGone
        connectivity[connect2, node2] = linkGone
    end

    # Remove the last link from the arrays since linkGone was either replaced by lastLink above, or already was the last link.
    links[:, lastLink] .= 0
    slipPlane[:, lastLink] .= 0
    bVec[:, lastLink] .= 0
    segForce[:, :, lastLink] .= 0
    network.numSeg[1] -= 1
    linksConnect[:, lastLink] .= 0

    return network
end

"""
```
mergeNode!(network::DislocationNetwork, nodeKept, nodeGone)
```
Merges `nodeGone` into `nodeKept`.
"""
function mergeNode!(network::DislocationNetwork, nodeKept, nodeGone)
    @assert nodeKept <= network.numNode[1] && nodeGone <= network.numNode[1] "mergeNode: the node kept after merging, $nodeKept and node removed after merging, $nodeGone, must be in the simulation."

    # Return if both nodes to be merged are the same.
    nodeKept == nodeGone && return 0, network

    elemT = eltype(network.bVec)

    nodeKeptConnect = network.connectivity[1, nodeKept]
    nodeGoneConnect = network.connectivity[1, nodeGone]
    totalConnect = nodeKeptConnect + nodeGoneConnect

    # Pass connections from nodeGone to nodeKept.
    if size(network.connectivity, 1) < 2 * totalConnect + 1
        network = DislocationNetwork(;
            links = network.links,
            slipPlane = network.slipPlane,
            bVec = network.bVec,
            coord = network.coord,
            label = network.label,
            segForce = network.segForce,
            nodeVel = network.nodeVel,
            nodeForce = network.nodeForce,
            numNode = network.numNode,
            numSeg = network.numSeg,
            maxConnect = network.maxConnect,
            linksConnect = network.linksConnect,
            connectivity = vcat(
                network.connectivity,
                zeros(
                    Int,
                    2 * totalConnect + 1 - size(network.connectivity, 1),
                    length(network.label),
                ),
            ),
            segIdx = network.segIdx,
        )
    end

    links = network.links
    slipPlane = network.slipPlane
    bVec = network.bVec
    connectivity = network.connectivity
    linksConnect = network.linksConnect

    # This next loop is equivalent to the following expression but doesn't allocate memory.
    # connectivity[(2 * (nodeKeptConnect + 1)):(2 * totalConnect + 1), nodeKept] = connectivity[2:(2 * nodeGoneConnect + 1), nodeGone]
    idx = (2 * (nodeKeptConnect + 1)):(2 * totalConnect + 1)
    lhs = idx[1]
    @inbounds @simd for i in 0:(length(idx) - 1)
        connectivity[lhs + i, nodeKept] = connectivity[2 + i, nodeGone]
    end
    connectivity[1, nodeKept] = totalConnect

    # Replace nodeGone with nodeKept in links and update linksConnect with the new positions of the links in connectivity.
    @inbounds @simd for i in 1:nodeGoneConnect
        link = connectivity[2 * i, nodeGone]      # Link id for nodeGone
        colLink = connectivity[2 * i + 1, nodeGone] # Position of links where link appears.
        links[colLink, link] = nodeKept         # Replace nodeGone with nodeKept.
        linksConnect[colLink, link] = nodeKeptConnect + i # Increase connectivity of nodeKept.
    end

    # Remove node from network and update index of nodeKept in case it changed.
    lastNode = maximum((network.numNode[1], 1))
    removeNode!(network, nodeGone, lastNode)
    coord = network.coord
    nodeVel = network.nodeVel
    connectivity = network.connectivity
    nodeKept == lastNode ? nodeKept = nodeGone : nothing

    # Warning: connectivity can be updated by removeLink!().
    # Delete self links of nodeKept.
    i = 1
    @inbounds while i < connectivity[1, nodeKept]
        link, missing, nodeNotLink = findConnectedNode(network, nodeKept, i) # i connection.

        # If nodeKept and nodeNotLink are the same, this is a self-link, therefore delete. Else increase counter.
        nodeKept == nodeNotLink ? removeLink!(network, link) : i += 1
        links = network.links
        connectivity = network.connectivity
    end

    # Warning: connectivity can be updated in the while loop.
    # Delete duplicate links.
    i = 1
    @inbounds while i < connectivity[1, nodeKept]
        link1, colNotLink1, nodeNotLink1 = findConnectedNode(network, nodeKept, i) # i connection.

        # Same but for the next link.
        j = i + 1
        while j <= connectivity[1, nodeKept]
            link2, colNotLink2, nodeNotLink2 = findConnectedNode(network, nodeKept, j) # i connection.

            # Continue to next iteration if no duplicate links are found.
            if nodeNotLink1 != nodeNotLink2
                j += 1
                continue
            end

            # Fix Burgers vector.
            # If the nodes are on different ends of a link, conservation of Burgers vector in a loop requires we subtract contributions from the two links involved. Else we add them.
            if colNotLink1 + colNotLink2 == 3
                @inbounds @simd for i in 1:3
                    bVec[i, link1] -= bVec[i, link2]
                end
            else
                @inbounds @simd for i in 1:3
                    bVec[i, link1] += bVec[i, link2]
                end
            end

            # WARNING This calculation is odd. Try using the cross product of the adjacent segments.
            # Fix slip plane.
            # Line direction and velocity of the resultant dislocation.
            t = SVector{3, elemT}(
                coord[1, nodeKept] - coord[1, nodeNotLink1],
                coord[2, nodeKept] - coord[2, nodeNotLink1],
                coord[3, nodeKept] - coord[3, nodeNotLink1],
            )

            v = SVector{3, elemT}(
                nodeVel[1, nodeKept] + nodeVel[1, nodeNotLink1],
                nodeVel[2, nodeKept] + nodeVel[2, nodeNotLink1],
                nodeVel[3, nodeKept] + nodeVel[3, nodeNotLink1],
            )

            # Burgers vector and potential new slip plane.
            b = SVector{3, elemT}(bVec[1, link1], bVec[2, link1], bVec[3, link1])
            n1 = t × b  # For non-screw segments.
            n2 = t × v  # For screw segments.
            if n1 ⋅ n1 > eps(elemT) # non-screw
                slipPlane[:, link1] = n1 / norm(n1)
            elseif n2 ⋅ n2 > eps(elemT) # screw
                slipPlane[:, link1] = n2 / norm(n2)
            end

            # Remove link2 from network and update the index of link1 in case it changed.
            lastLink = maximum((network.numSeg[1], 1))
            removeLink!(network, link2, lastLink)
            links = network.links
            connectivity = network.connectivity
            link1 == lastLink ? link1 = link2 : nothing

            # If the burgers vector of the new junction is non-zero, continue to the next iteration. Else remove it.
            b = SVector{3, elemT}(bVec[1, link1], bVec[2, link1], bVec[3, link1])
            if b ⋅ b ≈ 0
                removeLink!(network, link1)
                links = network.links
                connectivity = network.connectivity
                # If the node that was connected to nodeKept has no connections, remove it from the network and update the index of nodeKept in case it changed.
                if connectivity[1, nodeNotLink1] == 0
                    lastNode = maximum((network.numNode[1], 1))
                    removeNode!(network, nodeNotLink1, lastNode)
                    coord = network.coord
                    nodeVel = network.nodeVel
                    connectivity = network.connectivity
                    nodeKept == lastNode ? nodeKept = nodeNotLink1 : nothing
                end
            end

            # Since there was a change in the configuration, we check for duplicates again.
            i -= 1
            break
        end
        i += 1
    end

    # If nodeKept has no connections remove it.
    if connectivity[1, nodeKept] == 0
        removeNode!(network, nodeKept)
        nodeKept = 0
    end

    return nodeKept, network
end

"""
```
coarsenNetwork!(
    dlnParams::DislocationParameters,
    matParams::MaterialParameters,
    mesh::AbstractMesh,
    forceDisplacement::ForceDisplacement,
    network::DislocationNetwork,
)
```
Coarsens network such that no links are smaller than the minimum allowable length and so that no two links form triangles with area under the minimum allowed.
"""
function coarsenNetwork!(
    dlnParams::DislocationParameters,
    matParams::MaterialParameters,
    mesh::AbstractMesh,
    forceDisplacement::ForceDisplacement,
    network::DislocationNetwork,
)
    minAreaSq = dlnParams.minAreaSq
    minSegLen = dlnParams.minSegLen
    maxSegLen = dlnParams.maxSegLen

    label = network.label
    links = network.links
    coord = network.coord
    numNode = network.numNode[1]
    nodeVel = network.nodeVel
    connectivity = network.connectivity
    elemT = eltype(network.coord)

    i = 1
    while i <= numNode
        # We only want to coarsen real internal nodes. Else we skip to the next node.
        if !(connectivity[1, i] == 2 && label[i] == intMobDln)
            i += 1
            continue
        end

        missing, missing, link1_nodeOppI = findConnectedNode(network, i, 1) # 1st connection.
        missing, missing, link2_nodeOppI = findConnectedNode(network, i, 2) # 2nd connection.

        # We don't want to remesh out segments between two fixed nodes because the nodes by definition do not move and act as a source, thus we skip to the next node.
        if label[link1_nodeOppI] == intFixDln && label[link2_nodeOppI] == intFixDln
            i += 1
            continue
        end

        # Coordinate of node i
        iCoord = SVector{3, elemT}(coord[1, i], coord[2, i], coord[3, i])
        # Create a triangle formed by the three nodes involved in coarsening.
        coordVec1 =
            SVector{3, elemT}(
                coord[1, link1_nodeOppI],
                coord[2, link1_nodeOppI],
                coord[3, link1_nodeOppI],
            ) - iCoord # Vector between node 1 and the node it's connected to via link 1.

        coordVec2 =
            SVector{3, elemT}(
                coord[1, link2_nodeOppI],
                coord[2, link2_nodeOppI],
                coord[3, link2_nodeOppI],
            ) - iCoord # Vector between node 1 and the node it's connected to via link 2.
        coordVec3 = coordVec2 - coordVec1 # Vector between both nodes connected to iCoord.
        # Lengths of the triangle sides.
        r1 = norm(coordVec1)
        r2 = norm(coordVec2)
        r3 = norm(coordVec3)

        # If coarsening would result in a link whose length is bigger than the maximum allowed, and there r1 and r2 are bigger than the minimum allowed length, we skip to the next node.
        if r3 > maxSegLen && r1 > minSegLen && r2 > minSegLen
            i += 1
            continue
        end

        # Half the triangle's perimeter.
        r0 = (r1 + r2 + r3) / 2
        # Heron's formula for the area of a triangle when we know its sides.
        areaSq = r0 * (r0 - r1) * (r0 - r2) * (r0 - r3)

        # Node i velocities.
        iVel = SVector{3, elemT}(nodeVel[1, i], nodeVel[2, i], nodeVel[3, i])
        velVec1 =
            SVector{3, elemT}(
                nodeVel[1, link1_nodeOppI],
                nodeVel[2, link1_nodeOppI],
                nodeVel[3, link1_nodeOppI],
            ) - iVel
        velVec2 =
            SVector{3, elemT}(
                nodeVel[1, link2_nodeOppI],
                nodeVel[2, link2_nodeOppI],
                nodeVel[3, link2_nodeOppI],
            ) - iVel
        velVec3 = velVec2 - velVec1

        # Rate of change of side length with respect to time. We add eps(typeof(r)) to avoid division by zero.
        dr1dt = coordVec1 ⋅ velVec1 / (r1 + eps(elemT))
        dr2dt = coordVec2 ⋅ velVec2 / (r2 + eps(elemT))
        dr3dt = coordVec3 ⋅ velVec3 / (r3 + eps(elemT))
        dr0dt = (dr1dt + dr2dt + dr3dt) / 2

        # Rate of change of triangle area with respect to time.
        dAreaSqdt = dr0dt * (r0 - r1) * (r0 - r2) * (r0 - r3)
        dAreaSqdt += r0 * (dr0dt - dr1dt) * (r0 - r2) * (r0 - r3)
        dAreaSqdt += r0 * (r0 - r1) * (dr0dt - dr2dt) * (r0 - r3)
        dAreaSqdt += r0 * (r0 - r1) * (r0 - r2) * (dr0dt - dr3dt)

        # Coarsen if the area is less than the minimum allowed and shrinking. Or if the length of either of its links is less than the minimum allowed. Else continue to the next node.
        if !(areaSq < minAreaSq && dAreaSqdt < 0 || r1 < minSegLen || r2 < minSegLen)
            i += 1
            continue
        end

        # Merge node i into node link2_nodeOppI.
        nodeMerged, network = mergeNode!(network, link2_nodeOppI, i)
        getSegmentIdx!(network)
        links = network.links
        coord = network.coord
        label = network.label
        nodeVel = network.nodeVel
        numNode = network.numNode[1]
        connectivity = network.connectivity

        # If link2_nodeOppI no longer exists there is nothing to calculate and we proceed to the next iteration.
        if nodeMerged == 0
            i += 1
            continue
        end

        for j in 1:connectivity[1, nodeMerged]
            # Find the new link that has been created between nodeMerged and nodeNotMerged.
            linkMerged = connectivity[2 * j, nodeMerged]
            colLinkMerged = connectivity[2 * j + 1, nodeMerged]
            colNotLinkMerged = 3 - colLinkMerged
            nodeNotMerged = links[colNotLinkMerged, linkMerged]
            if nodeNotMerged == i || nodeNotMerged == link1_nodeOppI
                # Calculate segment force for segment linkMerged.
                calcSegForce!(
                    dlnParams,
                    matParams,
                    mesh,
                    forceDisplacement,
                    network,
                    linkMerged,
                )
                # Calculate node velocity.
                nodes = SVector{2, Int}(links[1, linkMerged], links[2, linkMerged])
                dlnMobility!(dlnParams, matParams, network, nodes)
                nodeVel = network.nodeVel
            end
        end
    end
    return network
end

max_stars_repo_name: dcelisgarza/DDD

max_forks_repo_name: dcelisgarza/DDD


text:
#########################################################################
#                                                                       #
# This file provides functions for interfacing with Matpower data files #
#                                                                       #
#########################################################################

"Parses the matpwer data from either a filename or an IO object"
function parse_matpower(io::IO; validate = true)::Dict
    mp_data = _parse_matpower_string(read(io, String))
    pm_data = _matpower_to_powermodels!(mp_data)
    if validate
        correct_network_data!(pm_data)
    end
    return pm_data
end

function parse_matpower(file::String; kwargs...)::Dict
    mp_data = open(file) do io
        parse_matpower(io; kwargs...)
    end
    return mp_data
end

### Data and functions specific to Matpower format ###

const _mp_data_names = [
    "mpc.version",
    "mpc.baseMVA",
    "mpc.bus",
    "mpc.gen",
    "mpc.branch",
    "mpc.dcline",
    "mpc.gencost",
    "mpc.dclinecost",
    "mpc.bus_name",
    "mpc.storage",
    "mpc.switch",
]

const _mp_bus_columns = [
    ("bus_i", Int),
    ("bus_type", Int),
    ("pd", Float64),
    ("qd", Float64),
    ("gs", Float64),
    ("bs", Float64),
    ("area", Int),
    ("vm", Float64),
    ("va", Float64),
    ("base_kv", Float64),
    ("zone", Int),
    ("vmax", Float64),
    ("vmin", Float64),
    ("lam_p", Float64),
    ("lam_q", Float64),
    ("mu_vmax", Float64),
    ("mu_vmin", Float64),
]

const _mp_bus_name_columns = [("name", Union{String, SubString{String}})]

const _mp_gen_columns = [
    ("gen_bus", Int),
    ("pg", Float64),
    ("qg", Float64),
    ("qmax", Float64),
    ("qmin", Float64),
    ("vg", Float64),
    ("mbase", Float64),
    ("gen_status", Int),
    ("pmax", Float64),
    ("pmin", Float64),
    ("pc1", Float64),
    ("pc2", Float64),
    ("qc1min", Float64),
    ("qc1max", Float64),
    ("qc2min", Float64),
    ("qc2max", Float64),
    ("ramp_agc", Float64),
    ("ramp_10", Float64),
    ("ramp_30", Float64),
    ("ramp_q", Float64),
    ("apf", Float64),
    ("mu_pmax", Float64),
    ("mu_pmin", Float64),
    ("mu_qmax", Float64),
    ("mu_qmin", Float64),
]

const _mp_branch_columns = [
    ("f_bus", Int),
    ("t_bus", Int),
    ("br_r", Float64),
    ("br_x", Float64),
    ("br_b", Float64),
    ("rate_a", Float64),
    ("rate_b", Float64),
    ("rate_c", Float64),
    ("tap", Float64),
    ("shift", Float64),
    ("br_status", Int),
    ("angmin", Float64),
    ("angmax", Float64),
    ("pf", Float64),
    ("qf", Float64),
    ("pt", Float64),
    ("qt", Float64),
    ("mu_sf", Float64),
    ("mu_st", Float64),
    ("mu_angmin", Float64),
    ("mu_angmax", Float64),
]

const _mp_dcline_columns = [
    ("f_bus", Int),
    ("t_bus", Int),
    ("br_status", Int),
    ("pf", Float64),
    ("pt", Float64),
    ("qf", Float64),
    ("qt", Float64),
    ("vf", Float64),
    ("vt", Float64),
    ("pmin", Float64),
    ("pmax", Float64),
    ("qminf", Float64),
    ("qmaxf", Float64),
    ("qmint", Float64),
    ("qmaxt", Float64),
    ("loss0", Float64),
    ("loss1", Float64),
    ("mu_pmin", Float64),
    ("mu_pmax", Float64),
    ("mu_qminf", Float64),
    ("mu_qmaxf", Float64),
    ("mu_qmint", Float64),
    ("mu_qmaxt", Float64),
]

const _mp_storage_columns = [
    ("storage_bus", Int),
    ("ps", Float64),
    ("qs", Float64),
    ("energy", Float64),
    ("energy_rating", Float64),
    ("charge_rating", Float64),
    ("discharge_rating", Float64),
    ("charge_efficiency", Float64),
    ("discharge_efficiency", Float64),
    ("thermal_rating", Float64),
    ("qmin", Float64),
    ("qmax", Float64),
    ("r", Float64),
    ("x", Float64),
    ("p_loss", Float64),
    ("q_loss", Float64),
    ("status", Int),
]

const _mp_switch_columns = [
    ("f_bus", Int),
    ("t_bus", Int),
    ("psw", Float64),
    ("qsw", Float64),
    ("state", Int),
    ("thermal_rating", Float64),
    ("status", Int),
]

""
function _parse_matpower_string(data_string::String)
    matlab_data, func_name, colnames = parse_matlab_string(data_string, extended = true)

    case = Dict{String, Any}()

    if func_name !== nothing
        case["name"] = func_name
    else
        @info(
            string(
                "no case name found in matpower file.  The file seems to be missing \"function mpc = ...\"",
            )
        )
        case["name"] = "no_name_found"
    end

    case["source_type"] = "matpower"
    if haskey(matlab_data, "mpc.version")
        case["source_version"] = matlab_data["mpc.version"]
    else
        @info(
            string(
                "no case version found in matpower file.  The file seems to be missing \"mpc.version = ...\"",
            )
        )
        case["source_version"] = "0.0.0+"
    end

    if haskey(matlab_data, "mpc.baseMVA")
        case["baseMVA"] = matlab_data["mpc.baseMVA"]
    else
        @info(
            string(
                "no baseMVA found in matpower file.  The file seems to be missing \"mpc.baseMVA = ...\"",
            )
        )
        case["baseMVA"] = 1.0
    end

    if haskey(matlab_data, "mpc.bus")
        buses = []
        for bus_row in matlab_data["mpc.bus"]
            bus_data = row_to_typed_dict(bus_row, _mp_bus_columns)
            bus_data["index"] = check_type(Int, bus_row[1])
            bus_data["source_id"] = ["bus", bus_data["index"]]
            push!(buses, bus_data)
        end
        case["bus"] = buses
    else
        error(
            string(
                "no bus table found in matpower file.  The file seems to be missing \"mpc.bus = [...];\"",
            ),
        )
    end

    if haskey(matlab_data, "mpc.gen")
        gens = []
        for (i, gen_row) in enumerate(matlab_data["mpc.gen"])
            gen_data = row_to_typed_dict(gen_row, _mp_gen_columns)
            gen_data["index"] = i
            gen_data["source_id"] = ["gen", i]
            push!(gens, gen_data)
        end
        case["gen"] = gens
    else
        error(
            string(
                "no gen table found in matpower file.  The file seems to be missing \"mpc.gen = [...];\"",
            ),
        )
    end

    if haskey(matlab_data, "mpc.branch")
        branches = []
        for (i, branch_row) in enumerate(matlab_data["mpc.branch"])
            branch_data = row_to_typed_dict(branch_row, _mp_branch_columns)
            branch_data["index"] = i
            branch_data["source_id"] = ["branch", i]
            push!(branches, branch_data)
        end
        case["branch"] = branches
    else
        error(
            string(
                "no branch table found in matpower file.  The file seems to be missing \"mpc.branch = [...];\"",
            ),
        )
    end

    if haskey(matlab_data, "mpc.dcline")
        dclines = []
        for (i, dcline_row) in enumerate(matlab_data["mpc.dcline"])
            dcline_data = row_to_typed_dict(dcline_row, _mp_dcline_columns)
            dcline_data["index"] = i
            dcline_data["source_id"] = ["dcline", i]
            push!(dclines, dcline_data)
        end
        case["dcline"] = dclines
    end

    if haskey(matlab_data, "mpc.storage")
        storage = []
        for (i, storage_row) in enumerate(matlab_data["mpc.storage"])
            storage_data = row_to_typed_dict(storage_row, _mp_storage_columns)
            storage_data["index"] = i
            storage_data["source_id"] = ["storage", i]
            push!(storage, storage_data)
        end
        case["storage"] = storage
    end

    if haskey(matlab_data, "mpc.switch")
        switch = []
        for (i, switch_row) in enumerate(matlab_data["mpc.switch"])
            switch_data = row_to_typed_dict(switch_row, _mp_switch_columns)
            switch_data["index"] = i
            switch_data["source_id"] = ["switch", i]
            push!(switch, switch_data)
        end
        case["switch"] = switch
    end

    if haskey(matlab_data, "mpc.bus_name")
        bus_names = []
        for (i, bus_name_row) in enumerate(matlab_data["mpc.bus_name"])
            bus_name_data = row_to_typed_dict(bus_name_row, _mp_bus_name_columns)
            bus_name_data["index"] = i
            bus_name_data["source_id"] = ["bus_name", i]
            push!(bus_names, bus_name_data)
        end
        case["bus_name"] = bus_names

        if length(case["bus_name"]) != length(case["bus"])
            error(
                "incorrect Matpower file, the number of bus names ($(length(case["bus_name"]))) is inconsistent with the number of buses ($(length(case["bus"]))).\n",
            )
        end
    end

    if haskey(matlab_data, "mpc.gencost")
        gencost = []
        for (i, gencost_row) in enumerate(matlab_data["mpc.gencost"])
            gencost_data = _mp_cost_data(gencost_row)
            gencost_data["index"] = i
            gencost_data["source_id"] = ["gencost", i]
            push!(gencost, gencost_data)
        end
        case["gencost"] = gencost

        if length(case["gencost"]) != length(case["gen"]) &&
           length(case["gencost"]) != 2 * length(case["gen"])
            error(
                "incorrect Matpower file, the number of generator cost functions ($(length(case["gencost"]))) is inconsistent with the number of generators ($(length(case["gen"]))).\n",
            )
        end
    end

    if haskey(matlab_data, "mpc.dclinecost")
        dclinecosts = []
        for (i, dclinecost_row) in enumerate(matlab_data["mpc.dclinecost"])
            dclinecost_data = _mp_cost_data(dclinecost_row)
            dclinecost_data["index"] = i
            dclinecost_data["source_id"] = ["dclinecost", i]
            push!(dclinecosts, dclinecost_data)
        end
        case["dclinecost"] = dclinecosts

        if length(case["dclinecost"]) != length(case["dcline"])
            error(
                "incorrect Matpower file, the number of dcline cost functions ($(length(case["dclinecost"]))) is inconsistent with the number of dclines ($(length(case["dcline"]))).\n",
            )
        end
    end

    for k in keys(matlab_data)
        if !in(k, _mp_data_names) && startswith(k, "mpc.")
            case_name = k[5:length(k)]
            value = matlab_data[k]
            if isa(value, Array)
                column_names = []
                if haskey(colnames, k)
                    column_names = colnames[k]
                end
                tbl = []
                for (i, row) in enumerate(matlab_data[k])
                    row_data = row_to_dict(row, column_names)
                    row_data["index"] = i
                    row_data["source_id"] = [case_name, i]
                    push!(tbl, row_data)
                end
                case[case_name] = tbl
                @info(
                    "extending matpower format with data: $(case_name) $(length(tbl))x$(length(tbl[1])-1)"
                )
            else
                case[case_name] = value
                @info("extending matpower format with constant data: $(case_name)")
            end
        end
    end

    return case
end

""
function _mp_cost_data(cost_row)
    ncost = cost_row[4]
    model = cost_row[1]
    if model == 1
        nr_parameters = ncost * 2
    elseif model == 2
        nr_parameters = ncost
    end

    cost_data = Dict(
        "model" => model,
        "startup" => check_type(Float64, cost_row[2]),
        "shutdown" => check_type(Float64, cost_row[3]),
        "ncost" => check_type(Int, cost_row[4]),
        "cost" => [check_type(Float64, x) for x in cost_row[5:(5 + nr_parameters - 1)]],
    )

    #=
    # skip this literal interpretation, as its hard to invert
    cost_values = [check_type(Float64, x) for x in cost_row[5:length(cost_row)]]
    if cost_data["model"] == 1:
        if length(cost_values)%2 != 0
            error("incorrect matpower file, odd number of pwl cost function values")
        end
        for i in 0:(length(cost_values)/2-1)
            p_idx = 1+2*i
            f_idx = 2+2*i
            cost_data["p_$(i)"] = cost_values[p_idx]
            cost_data["f_$(i)"] = cost_values[f_idx]
        end
    else:
        for (i,v) in enumerate(cost_values)
            cost_data["c_$(length(cost_values)+1-i)"] = v
        end
    =#
    return cost_data
end

### Data and functions specific to PowerModels format ###

"""
Converts a Matpower dict into a PowerModels dict
"""
function _matpower_to_powermodels!(mp_data::Dict{String, <:Any})
    pm_data = mp_data

    # required default values
    if !haskey(pm_data, "dcline")
        pm_data["dcline"] = []
    end
    if !haskey(pm_data, "gencost")
        pm_data["gencost"] = []
    end
    if !haskey(pm_data, "dclinecost")
        pm_data["dclinecost"] = []
    end
    if !haskey(pm_data, "storage")
        pm_data["storage"] = []
    end
    if !haskey(pm_data, "switch")
        pm_data["switch"] = []
    end

    # translate component models
    _mp2pm_branch!(pm_data)
    _mp2pm_dcline!(pm_data)

    # translate cost models
    _add_dcline_costs!(pm_data)

    # merge data tables
    _merge_bus_name_data!(pm_data)
    _merge_cost_data!(pm_data)
    _merge_generic_data!(pm_data)

    # split loads and shunts from buses
    _split_loads_shunts!(pm_data)

    # use once available
    arrays_to_dicts!(pm_data)

    for optional in ["dcline", "load", "shunt", "storage", "switch"]
        if length(pm_data[optional]) == 0
            pm_data[optional] = Dict{String, Any}()
        end
    end

    return pm_data
end

"""
    _split_loads_shunts!(data)

Seperates Loads and Shunts in `data` under separate "load" and "shunt" keys in the
PowerModels data format. Includes references to originating bus via "load_bus"
and "shunt_bus" keys, respectively.
"""
function _split_loads_shunts!(data::Dict{String, Any})
    data["load"] = []
    data["shunt"] = []

    load_num = 1
    shunt_num = 1
    for (i, bus) in enumerate(data["bus"])
        if bus["pd"] != 0.0 || bus["qd"] != 0.0
            append!(
                data["load"],
                [
                    Dict{String, Any}(
                        "pd" => bus["pd"],
                        "qd" => bus["qd"],
                        "load_bus" => bus["bus_i"],
                        "status" => convert(Int8, bus["bus_type"] != 4),
                        "index" => load_num,
                        "source_id" => ["bus", bus["bus_i"]],
                    ),
                ],
            )
            load_num += 1
        end

        if bus["gs"] != 0.0 || bus["bs"] != 0.0
            append!(
                data["shunt"],
                [
                    Dict{String, Any}(
                        "gs" => bus["gs"],
                        "bs" => bus["bs"],
                        "shunt_bus" => bus["bus_i"],
                        "status" => convert(Int8, bus["bus_type"] != 4),
                        "index" => shunt_num,
                        "source_id" => ["bus", bus["bus_i"]],
                    ),
                ],
            )
            shunt_num += 1
        end

        for key in ["pd", "qd", "gs", "bs"]
            delete!(bus, key)
        end
    end
end

"sets all branch transformer taps to 1.0, to simplify branch models"
function _mp2pm_branch!(data::Dict{String, Any})
    branches = [branch for branch in data["branch"]]
    if haskey(data, "ne_branch")
        append!(branches, data["ne_branch"])
    end
    for branch in branches
        if branch["tap"] == 0.0
            branch["transformer"] = false
            branch["tap"] = 1.0
        else
            branch["transformer"] = true
        end

        branch["g_fr"] = 0.0
        branch["g_to"] = 0.0

        branch["b_fr"] = branch["br_b"] / 2.0
        branch["b_to"] = branch["br_b"] / 2.0

        delete!(branch, "br_b")

        if branch["rate_a"] == 0.0
            delete!(branch, "rate_a")
        end
        if branch["rate_b"] == 0.0
            delete!(branch, "rate_b")
        end
        if branch["rate_c"] == 0.0
            delete!(branch, "rate_c")
        end
    end
end

"adds pmin and pmax values at to and from buses"
function _mp2pm_dcline!(data::Dict{String, Any})
    for dcline in data["dcline"]
        pmin = dcline["pmin"]
        pmax = dcline["pmax"]
        loss0 = dcline["loss0"]
        loss1 = dcline["loss1"]

        delete!(dcline, "pmin")
        delete!(dcline, "pmax")

        if pmin >= 0 && pmax >= 0
            pminf = pmin
            pmaxf = pmax
            pmint = loss0 - pmaxf * (1 - loss1)
            pmaxt = loss0 - pminf * (1 - loss1)
        end
        if pmin >= 0 && pmax < 0
            pminf = pmin
            pmint = pmax
            pmaxf = (-pmint + loss0) / (1 - loss1)
            pmaxt = loss0 - pminf * (1 - loss1)
        end
        if pmin < 0 && pmax >= 0
            pmaxt = -pmin
            pmaxf = pmax
            pminf = (-pmaxt + loss0) / (1 - loss1)
            pmint = loss0 - pmaxf * (1 - loss1)
        end
        if pmin < 0 && pmax < 0
            pmaxt = -pmin
            pmint = pmax
            pmaxf = (-pmint + loss0) / (1 - loss1)
            pminf = (-pmaxt + loss0) / (1 - loss1)
        end

        dcline["pmaxt"] = pmaxt
        dcline["pmint"] = pmint
        dcline["pmaxf"] = pmaxf
        dcline["pminf"] = pminf

        # preserve the old pmin and pmax values
        dcline["mp_pmin"] = pmin
        dcline["mp_pmax"] = pmax

        dcline["pt"] = -dcline["pt"] # matpower has opposite convention
        dcline["qf"] = -dcline["qf"] # matpower has opposite convention
        dcline["qt"] = -dcline["qt"] # matpower has opposite convention
    end
end

"adds dcline costs, if gen costs exist"
function _add_dcline_costs!(data::Dict{String, Any})
    if length(data["gencost"]) > 0 &&
       length(data["dclinecost"]) <= 0 &&
       length(data["dcline"]) > 0
        @info("added zero cost function data for dclines")
        model = data["gencost"][1]["model"]
        if model == 1
            for (i, dcline) in enumerate(data["dcline"])
                dclinecost = Dict(
                    "index" => i,
                    "model" => 1,
                    "startup" => 0.0,
                    "shutdown" => 0.0,
                    "ncost" => 2,
                    "cost" => [dcline["pminf"], 0.0, dcline["pmaxf"], 0.0],
                )
                push!(data["dclinecost"], dclinecost)
            end
        else
            for (i, dcline) in enumerate(data["dcline"])
                dclinecost = Dict(
                    "index" => i,
                    "model" => 2,
                    "startup" => 0.0,
                    "shutdown" => 0.0,
                    "ncost" => 3,
                    "cost" => [0.0, 0.0, 0.0],
                )
                push!(data["dclinecost"], dclinecost)
            end
        end
    end
end

"merges generator cost functions into generator data, if costs exist"
function _merge_cost_data!(data::Dict{String, Any})
    if haskey(data, "gencost")
        gen = data["gen"]
        gencost = data["gencost"]

        if length(gen) != length(gencost)
            if length(gencost) > length(gen)
                @warn(
                    "The last $(length(gencost) - length(gen)) generator cost records will be ignored due to too few generator records.",
                )
                gencost = gencost[1:length(gen)]
            else
                @warn(
                    "The number of generators ($(length(gen))) does not match the number of generator cost records ($(length(gencost))).",
                )
            end
        end

        for (i, gc) in enumerate(gencost)
            g = gen[i]
            @assert(g["index"] == gc["index"])
            delete!(gc, "index")
            delete!(gc, "source_id")

            _check_keys(g, keys(gc))
            merge!(g, gc)
        end

        delete!(data, "gencost")
    end

    if haskey(data, "dclinecost")
        dcline = data["dcline"]
        dclinecost = data["dclinecost"]

        if length(dcline) != length(dclinecost)
            @warn(
                "The number of dclines ($(length(dcline))) does not match the number of dcline cost records ($(length(dclinecost))).",
            )
        end

        for (i, dclc) in enumerate(dclinecost)
            dcl = dcline[i]
            @assert(dcl["index"] == dclc["index"])
            delete!(dclc, "index")
            delete!(dclc, "source_id")

            _check_keys(dcl, keys(dclc))
            merge!(dcl, dclc)
        end
        delete!(data, "dclinecost")
    end
end

"merges bus name data into buses, if names exist"
function _merge_bus_name_data!(data::Dict{String, Any})
    if haskey(data, "bus_name")
        # can assume same length is same as bus
        # this is validated during matpower parsing
        for (i, bus_name) in enumerate(data["bus_name"])
            bus = data["bus"][i]
            delete!(bus_name, "index")
            delete!(bus_name, "source_id")

            _check_keys(bus, keys(bus_name))
            merge!(bus, bus_name)
        end
        delete!(data, "bus_name")
    end
end

"merges Matpower tables based on the table extension syntax"
function _merge_generic_data!(data::Dict{String, Any})
    mp_matrix_names = [name[5:length(name)] for name in _mp_data_names]

    key_to_delete = []
    for (k, v) in data
        if isa(v, Array)
            for mp_name in mp_matrix_names
                if startswith(k, "$(mp_name)_")
                    mp_matrix = data[mp_name]
                    push!(key_to_delete, k)

                    if length(mp_matrix) != length(v)
                        error(
                            "failed to extend the matpower matrix \"$(mp_name)\" with the matrix \"$(k)\" because they do not have the same number of rows, $(length(mp_matrix)) and $(length(v)) respectively.",
                        )
                    end

                    @info(
                        "extending matpower format by appending matrix \"$(k)\" in to \"$(mp_name)\""
                    )

                    for (i, row) in enumerate(mp_matrix)
                        merge_row = v[i]
                        #@assert(row["index"] == merge_row["index"]) # note this does not hold for the bus table
                        delete!(merge_row, "index")
                        delete!(merge_row, "source_id")
                        for key in keys(merge_row)
                            if haskey(row, key)
                                error(
                                    "failed to extend the matpower matrix \"$(mp_name)\" with the matrix \"$(k)\" because they both share \"$(key)\" as a column name.",
                                )
                            end
                            row[key] = merge_row[key]
                        end
                    end

                    break # out of mp_matrix_names loop
                end
            end
        end
    end

    for key in key_to_delete
        delete!(data, key)
    end
end

""
function _check_keys(data, keys)
    for key in keys
        if haskey(data, key)
            error("attempting to overwrite value of $(key) in PowerModels data,\n$(data)")
        end
    end
end


text:
module SwitchTimeOpt

  # Import Necessary Modules
  using MathProgBase
  using DiffEqBase
  using OrdinaryDiffEq
  using Ipopt

  export stoproblem,
         setwarmstart!,setx0!,
         solve!,
         gettau,
         getdelta,
         gettaucomplete, getdeltacomplete,
         getobjval, getstat, getsoltime,
         getnobjeval, getngradeval, getnhesseval,
         simulatelinearized, simulateinput,
         simulate

  # Include Source Files
  include("types.jl")
  include("interface.jl")
  include("nlp.jl")
  include("simulation.jl")

end


text:
"""
Module MagmaThermoKinematics

Enables Earth Scientists to simulate the thermal evolution of magmatic systems.

"""
module MagmaThermoKinematics

# list required modules
using Random                                    # random numbers
using StructArrays                              # for tracers and dike polygon
using Parameters                                # More flexible definition of parameters
using Interpolations                            # Fast interpolations
using StaticArrays
#using CSV                                       # for reading phase diagrams

include("Units.jl")                             # various useful units

## Alphabetical include of computation-submodules (must be at end as needs to import from ParallelStencil).
include("Diffusion.jl")
export Diffusion2D, Diffusion3D  #

include("MeltingRelationships.jl")
export SolidFraction, ComputeLithostaticPressure, LoadPhaseDiagrams, PhaseDiagramData, ComputeDensityAndPressure
export PhaseRatioAverage!, ComputeSeismicVelocities, SolidFraction_Parameterized!

# Export functions that will be available outside this module
export StructArray, LazyRow # useful
export Tracer

include("Dikes.jl")
export Dike, DikePoly
export Tracer, AddDike, HostRockVelocityFromDike, CreatDikePolygon, volume_dike, InjectDike, TracersToGrid!

# routines related to advection & interpolation
include("Advection.jl")
export AdvectTemperature, Interpolate!, CorrectBounds, evaluate_interp_2D, evaluate_interp_3D

# Routines that deal with tracers
include("Tracers.jl")
export UpdateTracers, AdvectTracers!, InitializeTracers,PhaseRatioFromTracers, CorrectTracersForTopography!
export RockAssemblage

# Routines related to Parameters.jl, which come in handy in the main routine
export @unpack

end # module


text:
# BAI
# ===

# Binning index
typealias BinIndex Dict{UInt32,Vector{Chunk}}

# Linear index
typealias LinearIndex Vector{VirtualOffset}

const LinearWindowSize = 16 * 1024

type PseudoBin
    unmapped::Chunk
    n_mapped::UInt64
    n_unmapped::UInt64
end

type BAI
    indexes::Vector{Tuple{BinIndex,LinearIndex,Nullable{PseudoBin}}}
    n_no_coors::Nullable{UInt64}
end

function Base.read(input::IO, ::Type{BAI})
    # magic bytes
    B = read(input, UInt8)
    A = read(input, UInt8)
    I = read(input, UInt8)
    x = read(input, UInt8)
    if B != UInt8('B') || A != UInt8('A') || I != UInt8('I') || x != 0x01
        error("input is not a valid BAI file")
    end

    n_refs = read(input, Int32)
    indexes = read_indexes(input, n_refs)

    if !eof(input)
        n_no_coors = Nullable(read(input, UInt64))
    else
        n_no_coors = Nullable{UInt64}()
    end

    return BAI(indexes, n_no_coors)
end

# Read indexes for BAI and Tabix file formats.
function read_indexes(input, n_refs)
    indexes = Tuple{BinIndex,LinearIndex,Nullable{PseudoBin}}[]
    for _ in 1:n_refs
        # load a binning index (and a pseudo bin)
        n_bins = read(input, Int32)
        binindex = BinIndex()
        pbin = Nullable{PseudoBin}()
        for _ in 1:n_bins
            bin = read(input, UInt32)
            n_chunks = read(input, Int32)
            if bin == 37450
                # pseudo-bin
                @assert n_chunks == 2
                chunk_beg = read(input, UInt64)
                chunk_end = read(input, UInt64)
                n_mapped = read(input, UInt64)
                n_unmapped = read(input, UInt64)
                pbin = Nullable(PseudoBin(
                    Chunk(chunk_beg, chunk_end),
                    n_mapped,
                    n_unmapped))
            else
                chunks = Chunk[]
                for i in 1:n_chunks
                    chunk_beg = read(input, UInt64)
                    chunk_end = read(input, UInt64)
                    push!(chunks, Chunk(chunk_beg, chunk_end))
                end
                binindex[bin] = chunks
            end
        end

        # load a linear index
        n_intvs = read(input, Int32)
        linindex = LinearIndex()
        for _ in 1:n_intvs
            push!(linindex, read(input, UInt64))
        end

        push!(indexes, (binindex, linindex, pbin))
    end

    return indexes
end

# Calculate bins overlapping a region [from, to] (one-based).
function reg2bins(from, to)
    bins = UInt32[]
    bin_start = 0
    for scale in 29:-3:14
        for k in ((from - 1) >> scale):((to - 1) >> scale)
            push!(bins, bin_start + k)
        end
        bin_start = 8 * bin_start + 1
    end
    return bins
end


text:
#####
##### Annotation
#####

# Note that the real field type restrictions here are more lax than the documented
# ones for improved compatibility with data produced by older Onda.jl versions and/or
# non-Julia producers.
"""
    const Annotation = Legolas.@row("onda.annotation@1",
                                    recording::UUID,
                                    id::UUID,
                                    span::TimeSpan)

A type alias for [`Legolas.Row{typeof(Legolas.Schema("onda.annotation@1"))}`](https://beacon-biosignals.github.io/Legolas.jl/stable/#Legolas.@row)
representing an [`onda.annotation` as described by the Onda Format Specification](https://github.com/beacon-biosignals/Onda.jl##ondaannotation1).

This type primarily exists to aid in the validated row construction, and is not intended to be used
as a type constraint in function or struct definitions. Instead, you should generally duck-type any
"annotation-like" arguments/fields so that other generic row types will compose with your code.
"""
const Annotation = @row("onda.annotation@1",
                        recording::Union{UInt128,UUID} = UUID(recording),
                        id::Union{UInt128,UUID} = UUID(id),
                        span::Union{NamedTupleTimeSpan,TimeSpan} = TimeSpan(span))

"""
    write_annotations(io_or_path, table; kwargs...)

Invoke/return `Legolas.write(path_or_io, annotations, Schema("onda.annotation@1"); kwargs...)`.
"""
write_annotations(path_or_io, annotations; kwargs...) = Legolas.write(path_or_io, annotations, Legolas.Schema("onda.annotation@1"); kwargs...)

"""
    validate_annotations(annotations)

Perform both table-level and row-level validation checks on the content of `annotations`,
a presumed `onda.annotation` table. Returns `annotations`.

This function will throw an error in any of the following cases:

- `Legolas.validate(annotations, Legolas.Schema("onda.annotation@1"))` throws an error
- `Annotation(row)` errors for any `row` in `Tables.rows(annotations)`
- `annotations` contains rows with duplicate `id`s
"""
validate_annotations(annotations) = _fully_validate_legolas_table(annotations, Legolas.Schema("onda.annotation@1"), :id)

#####
##### utilities
#####

"""
    merge_overlapping_annotations(annotations)

Given the `onda.annotation`-compliant table `annotations`, return
a table corresponding to `annotations` except that overlapping entries have
been merged.

Specifically, two annotations `a` and `b` are determined to be "overlapping"
if `a.recording == b.recording && TimeSpans.overlaps(a.span, b.span)`. Merged
annotations' `span` fields are generated via calling `TimeSpans.shortest_timespan_containing`
on the overlapping set of source annotations.

The returned annotations table only has a single custom column named `from`
whose entries are `Vector{UUID}`s populated with the `id`s of the generated
annotations' source(s). Note that every annotation in the returned table
has a freshly generated `id` field and a non-empty `from` field, even if
the `from` only has a single element (i.e. corresponds to a single
non-overlapping annotation).

Note that this function internally works with `Tables.columns(annotations)`
rather than `annotations` directly, so it may be slower and/or require more
memory if `!Tables.columnaccess(annotations)`.
"""
function merge_overlapping_annotations(annotations)
    columns = Tables.columns(annotations)
    merged = Annotation[]
    for (rid, (locs,)) in Legolas.locations((columns.recording,))
        subset = (recording=view(columns.recording, locs), id=view(columns.id, locs), span=view(columns.span, locs))
        p = sortperm(subset.span; by=TimeSpans.start)
        sorted = Tables.rows((recording=view(subset.recording, p), id=view(subset.id, p), span=view(subset.span, p)))
        init = first(sorted)
        push!(merged, Annotation(recording=rid, id=uuid4(), span=init.span, from=[init.id]))
        for next in Iterators.drop(sorted, 1)
            prev = merged[end]
            if next.recording == prev.recording && TimeSpans.overlaps(next.span, prev.span)
                push!(prev.from, next.id)
                merged[end] = Annotation(Tables.rowmerge(prev; span=TimeSpans.shortest_timespan_containing((prev.span, next.span))))
            else
                push!(merged, Annotation(; recording=next.recording, id=uuid4(), span=next.span, from=[next.id]))
            end
        end
    end
    return merged
end


text:
#######################################
#
# sfp file
#
#######################################

@doc """
Read sfp file

#### Input
* `fname`: Name or path for the sfp file

#### Output
* `elec`: Electrodes object
""" ->
function read_sfp(fname::AbstractString; coordinate=Talairach)
    info("Reading dat file = $fname")

    # Create an empty electrode set
    elecs = Electrode[]

    # Read file
    df = readtable(fname, header = false, separator = ' ')

    # Save locations
    x = df[:x2]
    y = df[:x3]
    z = df[:x4]

    # Convert label to ascii and remove '
    labels = df[:x1]
    for i = 1:length(labels)
        push!(elecs, Electrode(replace(labels[i], "'", "" ), coordinate(x[i], y[i], z[i]), Dict()))
    end

    debug("Imported $(length(elecs)) electrodes")

    return elecs
end


text:
# Lifting Functions on type T to fnuctions on T-typed RandVars #

# No Exists{T} yet https://github.com/JuliaLang/julia/issues/21026#issuecomment-306624369"
function liftnoesc(fnm::Union{Symbol, Expr}, isrv::NTuple{N, Bool}) where N
  args = [isrv ?  :($(Symbol(:x, i))::Omega.RandVar) : Symbol(:x, i)  for (i, isrv) in enumerate(isrv)]
  quote
  function $fnm($(args...))
    Omega.lift($fnm)($(args...))
  end
  end
end

function liftesc(fnm::Union{Symbol, Expr}, isrv::NTuple{N, Bool}) where N
  args = [isrv ?  :($(Symbol(:x, i))::Omega.RandVar) : Symbol(:x, i)  for (i, isrv) in enumerate(isrv)]
  quote
  function $(esc(fnm))($(args...))
    Omega.lift($fnm)($(args...))
  end
  end
end

"Updates methods of `mod` to add lifted counterparts to function with name fnm"
function lift!(fnm::Union{Expr, Symbol}, n::Integer;
               mod::Module = @__MODULE__())
  # @show mod
  combs = rvcombinations(n)
  for comb in combs
    Core.eval(mod, liftnoesc(fnm, comb))
  end
end

function lift!(f; n=3, mod::Module = @__MODULE__())
  lift!(:($f), n; mod=mod)
end

# macro lift(fnm::Union{Symbol, Expr}, n::Integer)
#   combinations = Iterators.product(((true,false) for i = 1:n)...)
#   combinations = Iterators.filter(any, combinations)
#   Expr(:block, map(comb -> liftmacro(fnm, comb), combinations)...)
# end

"Combinations of RV or Not RV"
function rvcombinations(n)
  combinations = Iterators.product(((true,false) for i = 1:n)...)
  Iterators.filter(any, combinations)
end

# Pre Lifted #

fnms = [:(Base.:-),
        :(Base.:+),
        :(Base.:*),
        :(Base.:/),
        :(Base.:^),
        :(Base.:sin),
        :(Base.:cos),
        :(Base.:tan),
        :(Base.sum),
        :(Base.:&),
        :(Base.:|),
        :(Base.:sqrt),
        :(Base.:abs),
        :(Base.:exp),
        :(Base.:log),
        # :(Base.getindex),
        # :(Base.:(==)),
        :(Base.:>),
        :(Base.:>=),
        :(Base.:<=),
        :(Base.:<),
        :(Base.:!),
        ]

const MAXN = 4
for fnm in fnms, i = 1:MAXN
  lift!(fnm, i)
end

# lift(f::Function) = (args...) -> ciid(f, args...)

@inline liftciid(f, args...) = URandVar(liftreifyapply, (f, args...))
@inline liftreifyapply(ωπ, f, args...) = f(reify(ωπ, args)...)

@generated function maybelift(f, args...)
  if any([arg <: RandVar for arg in args])
    :(liftciid(f, args...))
  else
    :(f(args...))
  end
end

lift(f) = (args...) -> maybelift(f, args...)

# Special Cases $
"""Lifted equality:
x ==ᵣ y results in a Boolean valued random variable which asks is the
__realization__ (hence subscript r) of `x` equal to `y`
If either `x` (or `y`) is a constant (not a RandVar) then it determines if
realization of `x` (or `y`) is equal to y (or `x`)"""
x ==ᵣ y = lift(==)(x, y)


text:
# GitHub markdown
# ---------------

Base.@kwdef mutable struct GitHubMarkdown <: WeaveFormat
    description = "GitHub Markdown"
    extension = "md"
    codestart = "```julia"
    codeend = "```\n"
    termstart = codestart
    termend = codeend
    outputstart = "```"
    outputend = "```\n\n"
    fig_ext = ".png"
    mimetypes = ["image/png", "image/svg+xml", "image/jpg",
                "text/markdown", "text/plain"]
    out_width = nothing
    out_height = nothing
    fig_pos = nothing
    fig_env = nothing
    # specials
    preserve_header = true
end
register_format!("github", GitHubMarkdown())

function render_figures(docformat::GitHubMarkdown, chunk)
    fignames = chunk.figures
    caption = chunk.options[:fig_cap]
    result = ""
    figstring = ""

    length(fignames) > 0 || (return "")

    if !isnothing(caption)
        result *= "![$caption]($(fignames[1]))\n"
        for fig in fignames[2:end]
            result *= "![]($fig)\n"
            println("Warning, only the first figure gets a caption\n")
        end
    else
        for fig in fignames
            result *= "![]($fig)\n"
        end
    end
    return result
end

# Hugo markdown
# -------------

Base.@kwdef mutable struct Hugo <: WeaveFormat
    description = "Hugo Markdown (using shortcodes)"
    extension = "md"
    codestart = "```julia"
    codeend = "```\n"
    termstart = codestart
    termend = codeend
    outputstart = "```"
    outputend = "```\n\n"
    mimetypes = default_mime_types
    fig_ext = ".png"
    out_width = nothing
    out_height = nothing
    fig_pos = nothing
    fig_env = nothing
    # specials
    preserve_header = true
    uglyURLs = false # if `false`, prepend figure path by `..`
end
register_format!("hugo", Hugo())

function render_figures(docformat::Hugo, chunk)
    relpath = docformat.uglyURLs ? "" : ".."
    mapreduce(*, enumerate(chunk.figures), init = "") do (index, fig)
        if index > 1
            @warn("Only the first figure gets a caption.")
            title_spec = ""
        else
            caption = chunk.options[:fig_cap]
            title_spec = isnothing(caption) ? "" : "title=\"$(caption)\" "
        end
        "{{< figure src=\"$(joinpath(relpath, fig))\" $(title_spec) >}}"
    end
end

# multi language markdown
# -----------------------

Base.@kwdef mutable struct MultiMarkdown <: WeaveFormat
    description = "MultiMarkdown"
    extension = "md"
    codestart = "```julia"
    codeend = "```\n"
    termstart = codestart
    termend = codeend
    outputstart = "```"
    outputend = "```\n\n"
    mimetypes = default_mime_types
    fig_ext = ".png"
    out_width = nothing
    out_height = nothing
    fig_pos = nothing
    fig_env = nothing
    # specials
    preserve_header = true
end
register_format!("multimarkdown", MultiMarkdown())

function render_figures(docformat::MultiMarkdown, chunk)
    fignames = chunk.figures
    caption = chunk.options[:fig_cap]
    result = ""
    figstring = ""

    if chunk.options[:out_width] == nothing
        width = ""
    else
        width = "width=$(chunk.options[:out_width])"
    end

    length(fignames) > 0 || (return "")

    if !isnothing(caption)
        result *= "![$caption][$(fignames[1])]\n\n"
        result *= "[$(fignames[1])]: $(fignames[1]) $width\n"
        for fig in fignames[2:end]
            result *= "![][$fig]\n\n"
            result *= "[$fig]: $fig $width\n"
            println("Warning, only the first figure gets a caption\n")
        end
    else
        for fig in fignames
            result *= "![][$fig]\n\n"
            result *= "[$fig]: $fig $width\n"
        end
    end
    return result
end

# Rest
# ----

Base.@kwdef mutable struct Rest <: WeaveFormat
    description = "reStructuredText and Sphinx"
    extension = "rst"
    codestart = ".. code-block:: julia\n"
    codeend = "\n"
    termstart = codestart
    termend = codeend
    outputstart = "::\n"
    outputend = "\n\n"
    mimetypes = default_mime_types
    fig_ext = ".png"
    out_width = "15 cm"
    out_height = nothing
    fig_pos = nothing
    fig_env = nothing
    # specials
    indent = 4
end
register_format!("rst", Rest())

function render_figures(docformat::Rest, chunk)
    fignames = chunk.figures
    caption = chunk.options[:fig_cap]
    width = chunk.options[:out_width]
    result = ""
    figstring = ""

    for fig in fignames
        figstring *= @sprintf(".. image:: %s\n   :width: %s\n\n", fig, width)
    end

    if !isnothing(caption)
        result *= string(
            ".. figure:: $(fignames[1])\n",
            "   :width: $width\n\n",
            "   $caption\n\n",
        )
    else
        result *= figstring
        return result
    end
end

# Ansii
# -----

# asciidoc -b html5 -a source-highlighter=pygments ...
Base.@kwdef mutable struct AsciiDoc <: WeaveFormat
    description = "AsciiDoc"
    extension = "txt"
    codestart = "[source,julia]\n--------------------------------------"
    codeend = "--------------------------------------\n"
    termstart = codestart
    termend = codeend
    outputstart = "--------------------------------------"
    outputend = "--------------------------------------\n\n"
    mimetypes = default_mime_types
    fig_ext = ".png"
    out_width = "600"
    out_height = nothing
    fig_pos = nothing
    fig_env = nothing
end
register_format!("asciidoc", AsciiDoc())

function render_figures(docformat::AsciiDoc, chunk)
    fignames = chunk.figures
    caption = chunk.options[:fig_cap]
    width = chunk.options[:out_width]
    result = ""
    figstring = ""

    for fig in fignames
        figstring *= @sprintf("image::%s[width=%s]\n", fig, width)
    end

    if !isnothing(caption)
        result *= string("image::$(fignames[1])", "[width=$width,", "title=\"$caption\"]")
    else
        result *= figstring
        return result
    end
end