using Distributed
addprocs(2)

@everywhere using CharibdeOptim
@everywhere using ModelingToolkit
@everywhere using IntervalArithmetic

# --- Michalewicz Problem ---

@everywhere function michalewicz(X)
      sum = 0
      for i in 1:length(X)
            sum = sum + sin(X[i])*(sin(i*(X[i]^2)/π))^20
      end
      return -sum
end

(global_min, minimisers, information) = charibde_min(michalewicz, IntervalBox(0..π, 50))

# --- Sine Envolope Problem ---

@everywhere function sine_envolope(X)
      sum = 0
      for i in 1:length(X)-1
            temp = X[i+1]^2 + X[i]^2
            sum = sum + 0.5 + (sin(√temp -0.5)/(0.001*temp + 1))^2
      end
      return -sum
end

(global_min, minimisers, information) = charibde_min(sine_envolope, IntervalBox(-100..100, 5))

# --- Shekel Problem ---

@everywhere function shekel(X)
      sum = 0
      for i in 1:30
            s = 0
            for j in 1:length((X))
                  s = s + (X[j] - i*j)^2
            end
            sum = sum + (1 / (i + s))
      end
      return -sum
end

(global_min, minimisers, information) = charibde_min(shekel, IntervalBox(0..10, 5))

# --- Egg-Holder Problem ---

@everywhere function egg_holder(X)
      sum = 0
      for i in 1:length(X)-1
            sum = sum + ((X[i+1] + 47)sin(√abs(X[i+1] + 47 + X[i]/2)) + X[i]sin(√abs(X[i] - X[i+1] -47)))
      end
      return -sum
end

(global_min, minimisers, information) = charibde_min(egg_holder, IntervalBox(-512..512, 5))

# --- Rana Problem ---

@everywhere function rana(X)
      sum = 0
      for i in 1:length(X)-1
            temp1 = √abs(X[i+1] + X[i] + 1)
            temp2 = √abs(X[i+1] - X[i] + 1)
            sum = sum + (X[i]*cos(temp1) * sin(temp2) + (1+ X[i+1]) * sin(temp1) * cos(temp2))
      end
      return -sum
end

(global_min, minimisers, information) = charibde_min(rana , IntervalBox(-512..512, 5))

# --- Keane Problem ---

@everywhere function keane(X)
      temp_sum = 0
      temp_pro = 1
      temp_sum1 = 0
      for i in 1:length(X)
            temp_sum = temp_sum + cos(X[i])^4
            temp_pro = temp_pro * cos(X[i])^2
            temp_sum1 = temp_sum1 + i * X[i]^2
      end
      sum = abs(temp_sum - 2 * temp_pro) / √temp_sum1
      return -sum
end

@everywhere vars = ModelingToolkit.@variables x y w z
@everywhere C1 = constraint(vars , x*y*w*z, Interval(0.75, Inf))
@everywhere C2 = constraint(vars , x+y+z+w, Interval(-Inf, 30))

(global_min, minimisers, information) = charibde_min(keane, IntervalBox(0..10, 4), [C1, C2])


#############################################################################
# chambolle_pock.jl
# Uses the chambolle_pock algorithm to solve optimization problems consisting
# of two proxable parts f and g.
# For reference look at:
#
# function combi_split(problem, params = CombiParams())
#   @assert ...
#   combi_constraint = problem.combi_constraints[1]
#   z = zeros(size(problem.variable))
#   for t = 1:params.max_iters
#     # compute proximal operator of objective at z
#     reg_objective = problem.objective + 1/2*sumsquares(problem.variable - z)
#     conv_prob = minimize(reg_objective, problem.convex_constraints) # = prox_{objective}(z)
#     solve!(conv_prob)
#     x = evaluate(problem.variable)
#
#     # compute proximal operator of combi_constraint at x
#     z = prox!(combi_constraint, x)
#   end
#   fix!(problem.variable, z) # does this work?
#   problem.solution = z
#   return z
# end
#############################################################################

TOL = 1e-3

export chambolle_pock

function chambolle_pock(f::AbstractExpr, p::AssocPoly,
                        max_iters = 10000 :: Number)
  # Initialization
  x = get_cv(f)[1]
  x0 = zeros(x.size[1])
  u0 = zeros(x.size[1])
  para = Variable(x.size[1])

  # First iteration, in preparation of warm starts
  fix!(para, x0 - u0)
  convex_prob = Problem(:minimize, f + 0.5 * norm(x - para)^2)
  solve!(convex_prob, SCSSolver(verbose = false))
  x1 = (x.value)[:, 1]
  x0 = copy(x1)
  u0 = (u0 + 2 * x1 - x0) - prox(p, (u0 + 2 * x1 - x0), u0)

  # Step 2 iterations
  for iter = 1:max_iters
    fix!(para, x0 - u0)
    solve!(convex_prob, SCSSolver(verbose=false), warmstart = true)
    x1 = (x.value)[:, 1]
    u0 = (u0 + 2 * x1 - x0) - prox(p, (u0 + 2 * x1 - x0), u0)
    if sum(abs.(x1 - x0)) < TOL
      break
    end
    x0 = copy(x1)
  end
  x.value = x0
  return x0
end


"""
$(TYPEDEF)

A lorenz attractor with with two inputs and parameters `σ`, `ρ` and `β`.
"""
struct Lorenz3D{ΔT} <: ControlSystem{ΔT, 3, 2}
    σ::Float64;
    ρ::Float64;
    β::Float64;
end

function dx(cs::Lorenz3D, x::SVector{3}, u::SVector{2}, t::AbstractFloat)
    dx1 = cs.σ*(x[2] - x[1])
    dx2 = x[1]*(cs.ρ - x[3]) - x[2] + u[2]
    dx3 = x[1]*x[2] - cs.β*x[3] + u[1]

    return @SVector[dx1, dx2, dx3]
end


@symbol_func function beta_P(cur_reactor::AbstractReactor)
  cur_beta = 101325.0

  cur_beta *= 2

  cur_beta *= 4 * cur_reactor.pi * 1e-7

  cur_beta *= p_bar(cur_reactor)

  cur_beta /= B_P(cur_reactor) ^ 2

  cur_beta
end


function ueberpruefeMatrixVektor2x2(A,m)

Akorrekt = [1 0 1 0;
            0 1 0 1;
            1 1 0 0;
            0 0 1 1];

mkorrekt = [5 8  6 7];

# homogene Lösungen berechnen
homogenkorrekt = nullspace(Akorrekt);

# inhomogene Lösung berechnen
inhomogenkorrekt = pinv(Akorrekt)*mkorrekt';

homogen = nullspace(A);
inhomogen = pinv(A)*m';

u1 = homogen[:,1];
# u2 = homogen(:,2);
# u3 = homogen(:,3);
# u4 = homogen(:,4);
p = inhomogen;

Vergleich = Akorrekt'*mkorrekt';
VergleichSuS = A'*m';

if norm(Vergleich) == norm(VergleichSuS)
    println("Eure Matrix und euer Vektor sind korrekt!")
    println("Die Lösung eures Gleichungssystems lautet wie folgt, wobei r eine beliebige relle Zahl ist.")
    q = round.(p[:,1];digits=2);
    v1 = round.(u1[:,1];digits=2);
    print("Für f1 lautet die Lösung: ", q[1], " + r * ", v1[1],"\n")
    print("Für f2 lautet die Lösung: ", q[2], " + r * ", v1[2],"\n")
    print("Für f3 lautet die Lösung: ", q[3], " + r * ", v1[3],"\n")
    print("Für f4 lautet die Lösung: ", q[4], " + r * ", v1[4],"\n")
    #show(stdout,"text/plain",round.(q;digits=2))
    #show(stdout,"text/plain",round.(v1;digits=2))
        #println("p = $p")
    #println("u1 = $u1")
elseif norm(m) == norm([15 12 14 13])
    println("Formt euer Gleichungssystem um, so dass alle Variablen rechts und Zahlen links stehen. Stellt dann die Matrix-Vektorschreibweise auf.")
else
    println("Schaut euch noch einmal den Verlauf der Strahlen an und überdenkt eure Matrix.")
end

end


module FermiContour

module Transforms

include("JacobiFunc.jl")

import Base.(|>)

immutable Sn{M}
    m::M
end
@inline function |>(p, s::Sn)
    w,z = p
    sn,cn,dn = JacobiFunc.jacobi(z, s.m)
    return w*cn*dn, sn
end

immutable Möbius{A,B,C,D}
    a::A; b::B; c::C; d::D
end
@inline function |>(p, m::Möbius)
    w,z = p
    return w*(m.a*(m.c*z+m.d ) - m.c*(m.a*z+m.b))/(m.c*z+m.d)^2,
          (m.a*z+m.b)/(m.c*z+m.d)
end

immutable Affine{A,B}
    a::A; b::B
end
@inline function |>(p, a::Affine)
    w,z = p
    return w*a.a, a.a*z+a.b
end

immutable Sqrt end
@inline function |>(p, s::Sqrt)
    w,z = p
    return w/(2*sqrt(z)), sqrt(z)
end

end # module Transforms

include("JacobiFunc.jl")

export fermicontour
"""
    fermicontour(E1,E2,β,n) -> w,z

Compute a quadrature rule for evaluating the Fermi-Dirac function through contour integration.

For any `x ∈ [-E2,-E1] ∪ [E1,E2]` and any `β ∈ [0,∞]` it holds
```julia
fermidirac(x,β) ≈ sum(
    real(w*fermidirac(z,β)/(z-x))
    for (w,z) in zip(fermicontour(E1,E2,β,n)...)
)
```
and the error converges exponentially for `n → ∞`. Note that both `E1 = 0` (single interval)
and `β = Inf` (zero temperature) are valid inputs.

Reference: Lin Lin et al., Pole-Based Approximation of the Fermi-Dirac Function, Chinese Annals of Mathematics, Series B
"""
function fermicontour(E1,E2,β,n)
    m = E1^2 + π^2/β^2
    M = E2^2 + π^2/β^2

    k = (sqrt(M/m)-1)/(sqrt(M/m)+1)
    K = JacobiFunc.K(k^2)
    iK = im*JacobiFunc.iK(k^2)

    t = -K + iK/2 + 4*K*(0.5:1:2*n-0.5)/(2*n)
    w = Vector{Complex{typeof(k)}}(length(t))
    z = Vector{Complex{typeof(k)}}(length(t))
    for i = 1:length(t)
        # Quadrature points and weights
        w[i],z[i] = (-1/(2π*im) * 4*K/n, t[i]) |>
            Transforms.Sn(k^2) |>
            Transforms.Möbius(1,1/k, -1,1/k) |>
            Transforms.Affine(-sqrt(m*M), π^2/β^2) |>
            Transforms.Sqrt() |>
            Transforms.Affine(im,0)
    end
    return w,z
end

export fermidirac
"""
    fermidirac(z,β)

Compute `1/(1+exp(β*z))` in a numerically stable way.
"""
function fermidirac(z,β)
    return abs(β*z) < log(realmax(real(typeof(z)))) ?
            1/(1 + exp(β*z)) :
            0.5*(1-sign(real(z)))
end

end # module


using NSGAIII
using vOptGeneric, GLPKMathProgInterface

# using PyPlot

function plot_pop(P)
    clf()
    pop = filter(x -> x.CV ≈ 0, P)
    non_dom = NSGAIII.fast_non_dominated_sort!(P)[1]
    pop = setdiff(pop, non_dom)
    p = plot3D(map(x -> x.y[1], pop), map(x -> x.y[2], pop), map(x -> x.y[3], pop), "bo", markersize=1)
    p = plot3D(map(x -> x.y[1], non_dom), map(x -> x.y[2], non_dom), map(x -> x.y[3], non_dom), "go", markersize=1)
    ax = gca()
    ax[:set_xlim]([-95, 110])
    ax[:set_ylim]([-60, 45])
    ax[:set_zlim]([-50, 110])
    !isinteractive() && show()
    sleep(0.1)
end

m = vModel(solver = GLPKSolverMIP())

@variable(m, 0 <=x[1:5] <= 10)
@variable(m, δ[1:3], Bin)

@addobjective(m, Max, dot([17,-12,-12,-19,-6], x) + dot([-73, -99, -81], δ))
@addobjective(m, Max, dot([2,-6,0,-12,13], x) + dot([-61,-79,-53], δ))
@addobjective(m, Max, dot([-20,7,-16,0,-1], x) + dot([-72,-54,-79], δ))

@constraint(m, sum(δ) <= 1)
@constraint(m, -x[2] + 6x[5] + 25δ[1] <= 52)
@constraint(m, -x[1] + 18x[4] + 18x[5] + 8δ[2] <= 77)
@constraint(m, 7x[4] + 9x[5] + 19δ[3] <= 66)
@constraint(m, 16x[1] + 20x[5] <= 86)
@constraint(m, 13x[2] + 7x[4] <= 86)

print(m)

res = nsga(0, 100, 10, m,  pmut=0.3);

solve(m, method=:lex)
println("\nRésolution lexico-graphique : ")
display(getY_N(m));

println()
println("Meilleur individu sur le premier objectif")
x1 = sort(res, by = x -> x[2][1])[end];
println("x = $(x1[1][1:5]) , δ = $(x1[1][6:8])")
println("z = $(x1[2])")
println("CV : $(x1[3])")

println("Meilleur individu sur le deuxième objectif")
x2 = sort(res, by = x -> x[2][2])[end];
println("x = $(x2[1][1:5]) , δ = $(x2[1][6:8])")
println("z = $(x2[2])")
println("CV : $(x2[3])")

println("Meilleur individu sur le troisième objectif")
x3 = sort(res, by = x -> x[2][3])[end];
println("x = $(x3[1][1:5]) , δ = $(x3[1][6:8])")
println("z = $(x3[2])")
println("CV : $(x3[3])")


using MomentClosure
using MomentClosure: define_M, define_μ
using Symbolics: value, expand
using Test
using Catalyst

@parameters t, k_on, k_off, k_p, γ_p, b
@variables p(t), g(t)

vars = [g, p]
ps = [k_on, k_off, k_p, γ_p, b]
S = [1 -1 0 0;
   0 0 b -1]
as = [k_on*(1-g),  # 0 -> g
    k_off*g*(p^2), # g -> 0
    k_p*g,         # 0 -> mP, m ~ Geometric(mean=b)
    γ_p*p]         # p -> 0
binary_vars = [1]
rn = ReactionSystemMod(t, vars, ps, as, S)

μ = define_μ(2,5)
M = define_M(2,5)

sys = generate_raw_moment_eqs(rn, 3)
expr1 = sys.odes.eqs[4].rhs
expr2 = k_on*μ[0,1] + b*k_p*μ[2,0] - k_off*μ[1,3] - k_on*μ[1,1]-γ_p*μ[1,1]
@test isequal(expr1, expr2)
@test length(sys.odes.eqs) == 9

sys_clean = bernoulli_moment_eqs(sys, binary_vars)
expr1 = sys_clean.odes.eqs[3].rhs
expr2 = k_on*μ[0,1] + b*k_p*μ[1,0] - k_off*μ[1,3] - k_on*μ[1,1]-γ_p*μ[1,1]
expr2 = simplify(value.(expr2))
@test isequal(expr1, expr2)
@test length(sys_clean.odes.eqs) == 6

@test_throws ErrorException moment_closure(sys, "conditional derivative matching")

closed_eqs = moment_closure(sys, "conditional gaussian", binary_vars)
expr1 = closed_eqs.closure[μ[1,4]]
expr2 = 4*μ[1,1]*μ[1,3]*μ[1,0]^-1 + 3*μ[1,2]^2*μ[1,0]^-1 - 12*μ[1,2]*μ[1,1]^2*μ[1,0]^-2 + 6*μ[1,1]^4*μ[1,0]^-3
@test isequal(expr1, expr2)
expr1 = closed_eqs.closure[μ[1,3]]
expr2 = 3*μ[1,2]*μ[1,1]*μ[1,0]^-1 - 2*μ[1,1]^3*μ[1,0]^-2
@test isequal(expr1, expr2)

closed_eqs = moment_closure(sys, "conditional derivative matching", binary_vars)
expr1 = closed_eqs.closure[μ[1,4]]
expr2 = μ[1,3]^4*μ[1,1]^4*μ[1,0]^-1*μ[1,2]^-6
@test isequal(expr1, expr2)
expr1 = closed_eqs.closure[μ[1,3]]
expr2 = μ[1,2]^3*μ[1,0]*μ[1,1]^-3
@test isequal(expr1, expr2)

sys = generate_central_moment_eqs(rn, 3, 5)
expr1 = sys.odes.eqs[1].rhs
expr2 = -k_off*M[1,2] - k_on*μ[1,0] - k_off*M[0,2]*μ[1,0]-2*k_off*M[1,1]*μ[0,1] - k_off*μ[1,0]*μ[0,1]^2 + k_on
@test isequal(expand(expr1), expr2)
@test length(sys.odes.eqs) == 9

sys_clean = bernoulli_moment_eqs(sys, binary_vars)
expr1 = expand(sys_clean.odes.eqs[3].rhs)
expr2 = b*k_p*μ[1,0] - k_off*M[1,3] - k_on*M[1,1] - M[1,1]*γ_p - b*k_p*μ[1,0]^2 - k_off*M[0,3]*μ[1,0] -
      k_off*M[1,1]*μ[0,1]^2 - 2*k_off*M[1,2]*μ[0,1] - 2*k_off*M[0,2]*μ[0,1]*μ[1,0]
@test isequal(expr1, expr2)
@test length(sys_clean.odes.eqs) == 6

closed_eqs = moment_closure(sys, "conditional gaussian", binary_vars)
expr1 = expand(closed_eqs.closure[M[0,5]])
expr2 = 10*M[0,2]*M[0,3] + 5*M[0,4]*μ[0,1] - 15*μ[0,1]*M[0,2]^2
@test isequal(expr1, expr2)

closed_eqs = moment_closure(sys, "conditional derivative matching", binary_vars)
expr1 = closed_eqs.closure[M[1,3]]
expr2 = μ[1,0]*(M[1,1]+μ[0,1]*μ[1,0])^-3*(M[1,2]+M[0,2]*μ[1,0]+μ[1,0]*μ[0,1]^2+2*M[1,1]*μ[0,1])^3 -
    M[0,3]*μ[1,0] - 3*M[1,1]*μ[0,1]^2 - 3*M[1,2]*μ[0,1] - μ[1,0]*μ[0,1]^3- 3*M[0,2]*μ[0,1]*μ[1,0]
@test isequal(expand(expr1), expand(expr2))


using ITensors
using ITensorGLMakie
using Graphs

using ITensors.ContractionSequenceOptimization: optimal_contraction_sequence

N = 5
g = Graph(N)
g_edges = [2 => 3, 1 => 4, 1 => 5, 4 => 5]
for e in g_edges
  add_edge!(g, e)
end

A = itensornetwork(g; linkspaces=5)
sequence = optimal_contraction_sequence(A)
edge_labels = (; tags=true)
R = @visualize_sequence fig ITensors.contract(A; sequence=sequence) edge_labels =
  edge_labels

fig


using jInv.Mesh
using Multigrid
using SparseArrays
using LinearAlgebra
using Test
println("************************************************* GMG RAP for Elasticity 2D using Vanka ******************************************************");

domain = [0.0, 1.0, 0.0, 1.0];
n = [128,128];
Mr = getRegularMesh(domain,n)
mu = 1.0*ones(prod(Mr.n));
lambda = 10.0.*mu;
#Ar = GetLinearElasticityOperator(Mr,mu,lambda);
Ar = GetLinearElasticityOperatorMixedFormulation(Mr, mu,lambda);
println(norm(Ar - Ar'))
Ar = Ar + 1e-3*opnorm(Ar,1)*sparse(1.0I,size(Ar,2),size(Ar,2));

levels      = 4;
numCores 	= 2;
maxIter     = 5;
relativeTol = 1e-10;
relaxType   = "EconVankaFaces";
relaxParam  = 0.75;
relaxPre 	= 1;
relaxPost   = 1;
cycleType   ='V';
coarseSolveType = "Julia";
transferOperatorType = "SystemsFacesMixedLinear";

MG = getMGparam(Float64,Int64,levels,numCores,maxIter,relativeTol,relaxType,relaxParam,relaxPre,relaxPost,cycleType,coarseSolveType,0.5,0.0,transferOperatorType);

N = size(Ar,2);
b = Ar*rand(N); b = b/norm(b);
x = zeros(N);
MGsetup(Ar,Mr,MG,size(b,2),true);
println("****************************** Stand-alone GMG RAP: ******************************")
solveMG(MG,b,x,true);
@test norm(Ar*x - b) < 0.05;
println("****************************** GMG + CG: ******************************")
x[:].=0.0;
solveCG_MG(Ar,MG,b,x,true)
@test norm(Ar*x - b) < 0.01;

println("************************************************* GMG RAP for Elasticity 2D using HybridVanka(Kaczmarz) ******************************************************");

domain = [0.0, 1.0, 0.0, 1.0];
n = [128,128];
Mr = getRegularMesh(domain,n)
mu = 1.0*ones(prod(Mr.n));
lambda = 10.0.*mu;
#Ar = GetLinearElasticityOperator(Mr,mu,lambda);
Ar = GetLinearElasticityOperatorMixedFormulation(Mr, mu,lambda);
println(norm(Ar - Ar'))
Ar = Ar + 1e-3*opnorm(Ar,1)*sparse(1.0I,size(Ar,2),size(Ar,2));

levels      = 4;
numCores 	= 2;
maxIter     = 5;
relativeTol = 1e-10;
relaxType   = "hybridVankaFacesKaczmarz";
relaxParam  = getHybridKaczmarzParam([4,4],0.9,numCores,2);
relaxPre 	= 1;
relaxPost   = 1;
cycleType   ='V';
coarseSolveType = "Julia";
transferOperatorType = "SystemsFacesMixedLinear";

MG = getMGparam(Float64,Int64,levels,numCores,maxIter,relativeTol,relaxType,relaxParam,relaxPre,relaxPost,cycleType,coarseSolveType,0.5,0.0,transferOperatorType);

N = size(Ar,2);
b = Ar*rand(N); b = b/norm(b);
x = zeros(N);
MGsetup(Ar,Mr,MG,size(b,2),true);
println("****************************** Stand-alone GMG RAP: ******************************")
solveMG(MG,b,x,true);
@test norm(Ar*x - b) < 0.05;
println("****************************** GMG + CG: ******************************")
x[:].=0.0;
solveCG_MG(Ar,MG,b,x,true)
@test norm(Ar*x - b) < 0.01;

println("************************************************* GMG RAP for Elasticity 3D  using Vanka ******************************************************");

domain = [0.0, 1.0, 0.0, 1.0,0.0,1.0];
# for some reason this test fails to converge at [8,8,8] and succeeds at larger sizes. TODO: check this.
n = [8,8,8];
# n = [128,128,128];

Mr = getRegularMesh(domain,n)
mu = 1.0*ones(prod(Mr.n));
# lambda = 10.0.*mu;
lambda = 1.0.*mu;
Ar = GetLinearElasticityOperatorMixedFormulation(Mr, mu,lambda);
Ar = Ar + 1e-2*opnorm(Ar,1)*sparse(1.0I,size(Ar,2),size(Ar,2));
MG.relaxType = "VankaFaces";
MG.relaxParam = 0.5;

N = size(Ar,2);
b = Ar*rand(N); b = b/norm(b);
x = zeros(N);
clear!(MG)
MGsetup(Ar,Mr,MG,size(b,1),true);
println("****************************** Stand-alone GMG RAP: ******************************")
solveMG(MG,b,x,true);
@test norm(Ar*x - b) < 0.05;


lower_bound(::Type{MV}, Rₙ::AbstractVector{T},
            Rₙₖ::AbstractMatrix{T}) where {T<:AbstractFloat, MV<:MultivariateNormal} = sum(Rₙ)

"Gaussian Mixture Model"
function fit_mm(::Type{MV}, X::AbstractMatrix{T}, k::Int;
                tol::Real=1e-5,             # convergence tolerance
                maxiter::Integer=100,       # number of iterations
                μs::Union{AbstractArray{T,2}, Nothing} = nothing,
                Σs::Union{AbstractArray{T,3}, Nothing} = nothing,
                homoscedastic::Bool=false,
                init::Symbol=:kmeans,
                logprob::Bool=true,
                covreg::Real=1e-6,
                aitken::Bool=false
            ) where {T<:AbstractFloat, MV<:MultivariateNormal}

    d, n = size(X)
    Rₙ = zeros(T, n)
    Σ  = zeros(T, d, d)
    Z = similar(X)

    # initialize parameters
    πₖ, μₖ, Σₖ, Rₙₖ = initialize(MV, X, k, init=init)
    μs !== nothing && copyto!(μₖ, μs)
    Σs !== nothing && copyto!(Σₖ, Σs)
    Tmpₙ, TmpC = auxiliary(MV, X, k)

    ℒ = ℒ′ₐ = ℒₐ = ℒ′′ = ℒ′ = Δℒ = typemin(T)
    for itr in 1:maxiter

        # E Step: Calculate posterior probability
        #   Rₙₖ = E[ωₖ|x] = p(ωₖ=1|x) ∝ πₖ⋅p(x|ωₖ) = πₖ⋅(√(2π)^(-d/2)*|Σₖ⁻¹|)⋅exp(-0.5⋅(x-μₖ)ᵀΣₖ⁻¹(x-μₖ))
        # where ωₖ is the mixture indicator variable, s.t. ωₖ = 1 when the data point was generated by mixture ωₖ
        for j in 1:k
            # calculate responsibilities
            μ = μₖ[:,j]
            Σ = Σₖ[:,:,j]
            R = view(Rₙₖ,:,j)
            if logprob
                F = precision(MV, Σ)
                # logpdf!(R, μ, X, Hermitian(Σ), Z)
                logpdf!(R, X, μ, F, Z)
                R .+= log(πₖ[j])
            else
                pdf!(R, X, μ, Hermitian(Σ), Z)
                R .*= πₖ[j]
            end
        end
        # Calculate (log) responsibilities
        if logprob
            logsumexp!(Rₙ, Rₙₖ, Tmpₙ)
        else
            Rₙₖ[Rₙₖ .< eps(T)] .= eps(T)
            sum!(Rₙ, Rₙₖ)
            Rₙₖ ./= Rₙ
        end

        # Calculate lower bound value on the complete log-likelihood
        ℒ′′ = ℒ′
        ℒ′ = ℒ
        ℒ = lower_bound(MV, (logprob ? Rₙ : cliplog.(Rₙ)), Rₙₖ)

        # M Step: Calculate parameters
        stats!(πₖ, μₖ, X, Rₙₖ)
        covariances!(Σₖ, πₖ, μₖ, X, Rₙₖ, MV, covreg=covreg, aux=TmpC)
        πₖ ./= n

        if homoscedastic
            Σₖ[:,:,1] .*= πₖ[1]
            for j in 2:k
                Σⱼ = view(Σₖ,:,:,j)
                Σⱼ .*= πₖ[j]
                Σₖ[:,:,1] .+= Σⱼ
            end
            for j in 2:k
                Σₖ[:,:,j] .= copy(Σₖ[:,:,1])
            end
        end

        # Check convergence
        Δℒ = ℒ - ℒ′
        # aitken accelerated criterion
        if aitken
            a = Δℒ/(ℒ′-ℒ′′)
            ℒ′ₐ = ℒₐ
            ℒₐ = ℒ′+Δℒ/(1-a)
            Δℒ=ℒ′ₐ-ℒₐ
        end
        @debug "Likelihood" itr ℒ Δℒ
        abs(Δℒ) < tol && break
    end

    if Δℒ > tol
        @warn "No convergence" Δℒ=Δℒ tol=tol
    end

    return MixtureModel([distribution(MV, μₖ, Σₖ, j) for j in 1:k], πₖ)
end

distribution(::Type{FullNormal}, μₖ, Σₖ, j) = MvNormal(μₖ[:,j], Symmetric(Σₖ[:,:,j]))
distribution(::Type{DiagNormal}, μₖ, Σₖ, j) = MvNormal(μₖ[:,j], Diagonal(sqrt.(Σₖ[:,1,j])))
distribution(::Type{IsoNormal}, μₖ, Σₖ, j) = MvNormal(μₖ[:,j], sqrt(Σₖ[1,1,j]))


@testset "Maximum weight matching" begin
    g = complete_graph(3)
    w = [
        1 2 1
        1 1 1
        3 1 1
    ]
    match = maximum_weight_matching(g, w; optimizer=GLPK.Optimizer)
    @test match.mate[1] == 3
    @test match.weight ≈ 3

    g = complete_graph(3)
    w = zeros(3, 3)
    w[1, 2] = 1
    w[3, 2] = 1
    w[1, 3] = 1
    match = maximum_weight_matching(g, w; optimizer=GLPK.Optimizer)
    @test match.weight ≈ 1

    g = Graph(4)
    add_edge!(g, 1, 3)
    add_edge!(g, 1, 4)
    add_edge!(g, 2, 4)

    w = zeros(4, 4)
    w[1, 3] = 1
    w[1, 4] = 3
    w[2, 4] = 1

    match = maximum_weight_matching(g, w; optimizer=GLPK.Optimizer)
    @test match.weight ≈ 3
    @test match.mate[1] == 4
    @test match.mate[2] == -1
    @test match.mate[3] == -1
    @test match.mate[4] == 1

    g = Graph(4)
    add_edge!(g, 1, 2)
    add_edge!(g, 2, 3)
    add_edge!(g, 3, 1)
    add_edge!(g, 3, 4)
    match = maximum_weight_matching(g; optimizer=GLPK.Optimizer)
    @test match.weight ≈ 2
    @test match.mate[1] == 2
    @test match.mate[2] == 1
    @test match.mate[3] == 4
    @test match.mate[4] == 3

    w = zeros(4, 4)
    w[1, 2] = 1
    w[2, 3] = 1
    w[1, 3] = 1
    w[3, 4] = 1

    match = maximum_weight_matching(g, w; optimizer=GLPK.Optimizer)
    @test match.weight ≈ 2
    @test match.mate[1] == 2
    @test match.mate[2] == 1
    @test match.mate[3] == 4
    @test match.mate[4] == 3

    w = zeros(4, 4)
    w[1, 2] = 1
    w[2, 3] = 1
    w[1, 3] = 5
    w[3, 4] = 1

    match = maximum_weight_matching(g, w; optimizer=GLPK.Optimizer)
    @test match.weight ≈ 5
    @test match.mate[1] == 3
    @test match.mate[2] == -1
    @test match.mate[3] == 1
    @test match.mate[4] == -1
end

@testset "Maximum weight maximal matching" begin
    g = complete_bipartite_graph(2, 2)
    w = zeros(4, 4)
    w[1, 3] = 10.0
    w[1, 4] = 1.0
    w[2, 3] = 2.0
    w[2, 4] = 11.0
    match = maximum_weight_maximal_matching(g, w; optimizer=GLPK.Optimizer)
    @test match.weight ≈ 21
    @test match.mate[1] == 3
    @test match.mate[3] == 1
    @test match.mate[2] == 4
    @test match.mate[4] == 2

    g = complete_bipartite_graph(2, 4)
    w = zeros(6, 6)
    w[1, 3] = 10
    w[1, 4] = 0.5
    w[2, 3] = 11
    w[2, 4] = 1
    match = maximum_weight_maximal_matching(g, w; optimizer=GLPK.Optimizer)
    @test match.weight ≈ 11.5
    @test match.mate[1] == 4
    @test match.mate[4] == 1
    @test match.mate[2] == 3
    @test match.mate[3] == 2

    g = complete_bipartite_graph(2, 6)
    w = zeros(8, 8)
    w[1, 3] = 10
    w[1, 4] = 0.5
    w[2, 3] = 11
    w[2, 4] = 1
    w[2, 5] = -1
    w[2, 6] = -1
    match = maximum_weight_maximal_matching(g, w, 0; optimizer=GLPK.Optimizer)
    @test match.weight ≈ 11.5
    @test match.mate[1] == 4
    @test match.mate[4] == 1
    @test match.mate[2] == 3
    @test match.mate[3] == 2

    g = complete_bipartite_graph(4, 2)
    w = zeros(6, 6)
    w[3, 5] = 10
    w[3, 6] = 0.5
    w[2, 5] = 11
    w[1, 6] = 1
    w[1, 5] = -1

    match = maximum_weight_maximal_matching(g, w, 0; optimizer=GLPK.Optimizer)
    @test match.weight ≈ 12
    @test match.mate[1] == 6
    @test match.mate[2] == 5
    @test match.mate[3] == -1
    @test match.mate[4] == -1
    @test match.mate[5] == 2
    @test match.mate[6] == 1
end


using AbstractGPs, KernelFunctions, Distributions, Soss, LinearAlgebra, Plots, Random

# Adapted from a post by Arthur Lui
# https://discourse.julialang.org/t/gaussian-process-model-with-turing/42453/32?u=cscherrer

# Define model.
m = @model x begin
    # Priors.
    α ~ LogNormal(0.0, 0.1)
    ρ ~ LogNormal(0.0, 1.0)
    σ ~ LogNormal(0.0, 1.0)

    # Covariance function.
    kernel = α^2 * transform(SEKernel(), 1/(ρ*√2))

    # GP (implicit zero-mean).
    gp = GP(kernel)

    # Sampling Distribution (MvNormal likelihood).
    y ~ gp(x, σ^2 + 1e-6)  # add 1e-6 for numerical stability.
end

Random.seed!(1)
# Get some fake data
x = randn(20)

y =  sinpi.(x) .+ 0.1 .* randn(20);

# Sample from the posterior
post = dynamicHMC(m(x=x), (;y=y))

pred = [posterior(p._y_dist, y) for p in predict(Soss.withdistributions(m)(x=x), post)]

# gps = [posterior(p.gp(x), y) for p in pred]

plt = scatter(x, y; legend=false, dpi=200)

function plotgps(pred, x, y)
    plt = plot(;dpi=200)
    xx = collect(range(minimum(x)-0.5, maximum(x)+0.5; length=200))
    for p in pred
        try
            gp = posterior(p._y_dist, y)
            plot!(plt, xx, rand(gp(xx,1e-9),2); alpha=0.01, color="black", legend=false)
        catch
        end
    end
    scatter!(plt, x, y; color=1)
    ylims!(plt, -2.5, 2.5)
    return plt
end

plotgps(pred, x, y)

particles(post)

[posterior(p._y_dist, y) for p in predict(Soss.withdistributions(m)(x=x), post)]


# power analysis functions for mean for normal distribution with known variance

"""
    powerRTest(r::Real = 0, n::Real = 0, alpha::Float64 = 0.05, alternative = "two.sided")

Compute power of a sample of `n` observations to test the effect size `r`
at type I error = `alpha` (default: 0.05).
The effect size `r` is a linear correlation coefficient. The option
`alternative` is `two.sided` for H₀: r = 0 vs H₁: r ≠ 0, `less` for H₁: r < 0, and `greater`
for H₁: r > 0.
"""
function powerRTest(;
    r::Real = 0.0,
    n::Real = 0,
    alpha::Real = 0.05,
    alternative::String = "two")

    check_args(r=r,n=n,alpha=alpha)

    if alternative == "less"
        tside = 1
        r = -r
    elseif alternative in ("two","two.sided","two-sided","two sided")
        tside = 2
        r = abs(r)
    elseif alternative == "greater"
        tside = 1
    end

    if tside == 1
        ttt = cquantile(TDist(n-2),alpha)
        rc = sqrt(ttt^2 / (ttt^2 + n - 2))
        zr = atanh(r) + r / (2 * (n - 1))
        zrc = atanh(rc)
        return cdf(Normal(),(zr - zrc) * sqrt(n-3))
    end
    ttt = cquantile(TDist(n-2),alpha/2)
    rc = sqrt(ttt^2 / (ttt^2 + n - 2))
    zr = atanh(r) + r / (2 * (n - 1))
    zrc = atanh(rc)
    return cdf(Normal(),(zr - zrc) * sqrt(n-3)) + cdf(Normal(),(-zr - zrc) * sqrt(n-3))
end

"""
    samplesizeRTest(r::Real = 0, alpha::Float64 = 0.05, power::Float64 = 0.8, alternative = "two.sided")

Compute the sample size to detect the effect size `r` with power > `power`
at type I error = `alpha` (default: 0.05).
The effect size `r` is a linear correlation coefficient. The option
`alternative` is `two.sided` for H₀: r = 0 vs H₁: r ≠ 0, `less` for H₁: r < 0, and `greater`
for H₁: r > 0.
"""
function samplesizeRTest(;
    r::Real = 0.0,
    alpha::Real = 0.05,
    power::Real = 0.0,
    alternative::String = "two")

    check_args(r=r,alpha=alpha,power=power)

    return ceil(Int64,fzero(x->powerRTest(n = x, r = r, alpha = alpha, alternative = alternative) - power, 4.0+1e-10, 1e+09))
end

"""
    effectsizeRTest(n::Real = 0, alpha::Float64 = 0.05, power::Float64 = 0.8, alternative = "two.sided")

Compute the effect size that a sample size `n` can detect with power > `power`
at type I error = `alpha` (default: 0.05).
The effect size `r` is a linear correlation coefficient. The option
`alternative` is `two.sided` for H₀: r = 0 vs H₁: r ≠ 0, `less` for H₁: r < 0, and `greater`
for H₁: r > 0.
"""
function effectsizeRTest(;
    n::Real = 0,
    alpha::Real = 0.05,
    power::Real = 0.0,
    alternative::String = "two"
    )

    check_args(n=n,alpha=alpha,power=power)

    return fzero(x -> powerRTest(n = n, r = x, alpha = alpha, alternative = alternative) - power,.001,100)
end

"""
    alphaRTest(r::Real = 0, n::Real = 0, power::Float64 = 0.8, alternative = "two.sided")

Compute the probability of type I error that the sample size `n` has in detecting the
effect size `r` with power > `power` at type I error = `alpha` (default: 0.05).
The effect size `r` is a linear correlation coefficient. The option
`alternative` is `two.sided` for H₀: r = 0 vs H₁: r ≠ 0, `less` for H₁: r < 0, and `greater`
for H₁: r > 0.
"""
function alphaRTest(;
    n::Real = 0,
    r::Real = 0.0,
    power::Real = 0.0,
    alternative::String = "two"
    )

    check_args(r=r,n=n,power=power)

    return fzero(x->powerRTest(n = n, r = r, alpha = x, alternative = alternative) - power, 1e-10, 1 - 1e-10)
end

"""
    pwr.RTest(r::Real = 0, n::Real = 0, alpha::Float64 = 0.05, power::Float64 = 0.8, alternative = "two.sided")

Compute one of the test parameters such as sample size (`n`),
effect size (`r`), type I error (`alpha`), and power (`power`).
The parameter to be estimated must be set to zero (default).
The effect size `r` is a linear correlation coefficient. The option
`alternative` is `two.sided` for H₀: r = 0 vs H₁: r ≠ 0, `less` for H₁: r < 0, and `greater`
for H₁: r > 0.
"""
function RTest(;
    n::Real = 0,
    r::Real = 0.0,
    alpha::Float64 = 0.0,
    power::Float64 = 0.0,
    alternative::String = "two")

    if sum([x == 0 for x in (n,r,alpha,power)]) != 1
        error("exactly one of n, r, power, and alpha must be zero")
    end

    if power == 0.0
        power = powerRTest(n = n, r = r, alpha = alpha, alternative = alternative)
    elseif alpha == 0.0
        alpha = alphaRTest(n = n, r = r, power = power, alternative = alternative)
    elseif r == 0.0
        r = effectsizeRTest(n = n, alpha = alpha, power = power, alternative = alternative)
    elseif n == 0
        n = samplesizeRTest(r = r, alpha = alpha, power = power, alternative = alternative)
    end

    alt = Dict("two" => "two-sided","two.sided" => "two-sided", "two sided" => "two-sided", "less" => "less", "greater" => "greater")

    return htest(
        string("Approximate correlation power calculation (arctangh transformation)"),
        OrderedDict(
            "n" => n,
            "r" => r,
            "alpha" => alpha,
            "power" => power,
            "alternative" => alt[alternative])
        )
end


"""
    explicit_euler_riemann(mu, beta, p0, h, t0, t1, a0, a1)

Integrates a population without migration using Explicit Euler-Riemann.
mu is a mortality function of age and time. Beta is a birth rate of
age and time. p0 is an initial population, as a function of age.
h is the step size in both age and time. (t0, t1) are time limits.
(a0, a1) are age limits. This returns a single array in age and time
with the value of the population.

You must have h*mu < 1.
"""
function explicit_euler_riemann(mu, beta, p0, h, t0, t1, a0, a1)
    a = a0:h:a1
    t = t0:h:t1
    m = length(a)
    n = length(t)

    u = zeros(h, m, n)
    u[:, 1] = p0(a, t0)

    for i in 1:(n - 1)
        u[:, i + 1] = (1 .- h .* mu(a, t[i])) .* u[:, i]
        u[1, i + 1] = h * sum(beta(a, t[i]) .* u[1:m, i])
    end
    u
end

"""
    implicit_euler_riemann(mu, beta, p0, h, t0, t1, a0, a1)

Integrates a population without migration using Implicit Euler-Riemann.
mu is a mortality function of age and time. Beta is a birth rate of
age and time. p0 is an initial population, as a function of age.
h is the step size in both age and time. (t0, t1) are time limits.
(a0, a1) are age limits. This returns a single array in age and time
with the value of the population.

You must have h*mu < 1.
"""
function implicit_euler_riemann(mu, beta, p0, h, t0, t1, a0, a1)
    a = a0:h:a1
    t = t0:h:t1
    m = length(a)
    n = length(t)

    u = zeros(h, m, n)
    u[:, 1] = p0(a, t0)

    for i in 1:(n - 1)
        u[:, i + 1] = u[:, i] ./ (1 .+ h .* mu(a, t[i]))
        u[1, i + 1] = h * sum(beta(a, t[i]) .* u[1:m, i])
    end
    u
end

"""
Solve the problem for q(a, t) = p(a, t) / survival.
Then multiply the survival back in order to get p(a, t).
"""
function euler_riemann_survival(s, beta, p0, h, t0, t1, a0, a1)
    a = a0:h:a1
    t = t0:h:t1
    m = length(a)
    n = length(t)

    u = zeros(h, m, n)
    u[:, 1] = p0(a, t0)

    for i in 1:(n - 1)
        birth = sum(beta(a, t[i]) .* s[a, t[i]] .* u[1:m, i])
        for j in 1:min(m, n - i)
            u[j, i + j] = birth
        end
        u[:, i] .*= s[a, t[i]]
    end
    u[:, n] .*= s[a, t[n]]
    u
end


using Atum
using Atum.Advection

using Test
using Printf
using StaticArrays: SVector
using LinearAlgebra: norm

if !@isdefined integration_testing
  const integration_testing = parse(
    Bool,
    lowercase(get(ENV, "ATUM_INTEGRATION_TESTING", "false")),
  )
end

function wave(law, x⃗, t)
  ρ = 2 + sin(π * (sum(x⃗) - sum(constants(law).u⃗) * t))
  SVector(ρ)
end

function run(A, FT, N, K; volume_form=WeakForm())
  Nq = N + 1

  law = AdvectionLaw{FT, 1}()

  cell = LobattoCell{FT, A}(Nq)
  v1d = range(FT(-1), stop=FT(1), length=K+1)
  grid = brickgrid(cell, (v1d,); periodic=(true,))

  dg = DGSEM(; law, grid, volume_form,
               surface_numericalflux = RusanovFlux())

  cfl = FT(1 // 2)
  dt = cfl * step(v1d) / N / norm(constants(law).u⃗)
  timeend = FT(0.7)

  q = wave.(Ref(law), points(grid), FT(0))

  @info @sprintf """Starting
  N           = %d
  K           = %d
  volume_form = %s
  norm(q)     = %.16e
  """ N K volume_form weightednorm(dg, q)

  odesolver = LSRK54(dg, q, dt)
  solve!(q, timeend, odesolver)

  qexact = wave.(Ref(law), points(grid), timeend)
  errf = weightednorm(dg, q .- qexact)

  @info @sprintf """Finished
  norm(q)      = %.16e
  norm(q - qe) = %.16e
  """ weightednorm(dg, q) errf
  errf
end

let
  A = Array
  FT = Float64
  N = 4

  expected_error = Dict()

  #form, lev
  expected_error[WeakForm(), 1] = 1.5238393162339734e-04
  expected_error[WeakForm(), 2] = 4.8377348010293410e-06
  expected_error[WeakForm(), 3] = 1.4063471307765498e-07
  expected_error[WeakForm(), 4] = 4.5781066120767573e-09

  expected_error[FluxDifferencingForm(CentralFlux()), 1] = 1.5238393162391597e-04
  expected_error[FluxDifferencingForm(CentralFlux()), 2] = 4.8377348011318472e-06
  expected_error[FluxDifferencingForm(CentralFlux()), 3] = 1.4063471309426943e-07
  expected_error[FluxDifferencingForm(CentralFlux()), 4] = 4.5781064622952347e-09

  nlevels = integration_testing ? 4 : 1

  @testset for volume_form in (WeakForm(), FluxDifferencingForm(CentralFlux()))
    errors = zeros(FT, nlevels)
    for l in 1:nlevels
      K = 5 * 2 ^ (l - 1)
      errors[l] = run(A, FT, N, K; volume_form)
      @test errors[l] ≈ expected_error[volume_form, l]
    end
    if nlevels > 1
      rates = log2.(errors[1:(nlevels-1)] ./ errors[2:nlevels])
      @info "Convergence rates\n" *
        join(["rate for levels $l → $(l + 1) = $(rates[l])" for l in 1:(nlevels - 1)], "\n")
      @test rates[end] ≈ N + 1 atol=0.1
    end
  end
end


@testset "Evolution" begin
    @testset "Exponent function" begin
    @test evolve_operator(zeros(4, 4), 1.) ≈ Matrix{Float64}(I,4,4)
    A = [1. im 2;
    3 4 5im;
    6 7 8]
    @test evolve_operator(A,4.2) ≈ exp(4.2*A)

    # type check
    @test_throws MethodError evolve_operator(spzeros(4, 4), 1.)
  end

  @testset "Evolution functions" begin
    #size 4x4
    A = zeros(ComplexF64,4,4)
    A[2,2] = 1/2
    A[4,4] = 1/2
    S = [0.0+0.0im 0.5+0.0im 0.0+0.0im 0.0+0.0im 0.5+0.0im 0.5+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im;
        -0.5+0.0im -0.25+0.0im 0.0+0.5im 0.0+0.0im 0.0+0.0im 0.5+0.0im 0.5+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im;
        0.0+0.0im 0.0+0.5im -0.25+0.0im 0.0+0.5im 0.0+0.0im 0.0+0.0im 0.5+0.0im 0.5+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im;
        0.0+0.0im 0.0+0.0im 0.0+0.5im -0.25+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.5+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im;
        -0.5+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im -0.25+0.0im 0.5+0.0im 0.0+0.0im 0.0+0.0im 0.0-0.5im 0.5+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im;
        0.0+0.0im -0.5+0.0im 0.0+0.0im 0.0+0.0im -0.5+0.0im -0.5+0.0im 0.0+0.5im 0.0+0.0im 0.0+0.0im 0.0-0.5im 0.5+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im;
        0.0+0.0im 0.0+0.0im -0.5+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.5im -0.5+0.0im 0.0+0.5im 0.0+0.0im 0.0+0.0im 0.0-0.5im 0.5+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im;
        0.0+0.0im 0.0+0.0im 0.0+0.0im -0.5+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.5im -0.5+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0-0.5im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im;
        0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0-0.5im 0.0+0.0im 0.0+0.0im 0.0+0.0im -0.25+0.0im 0.5+0.0im 0.0+0.0im 0.0+0.0im 0.0-0.5im 0.5+0.0im 0.0+0.0im 0.0+0.0im;
        0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0-0.5im 0.0+0.0im 0.0+0.0im -0.5+0.0im -0.5+0.0im 0.0+0.5im 0.0+0.0im 0.0+0.0im 0.0-0.5im 0.5+0.0im 0.0+0.0im;
        0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0-0.5im 0.0+0.0im 0.0+0.0im 0.0+0.5im -0.5+0.0im 0.0+0.5im 0.0+0.0im 0.0+0.0im 0.0-0.5im 0.5+0.0im;
        0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0-0.5im 0.0+0.0im 0.0+0.0im 0.0+0.5im -0.5+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0-0.5im;
        0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0-0.5im 0.0+0.0im 0.0+0.0im 0.0+0.0im -0.25+0.0im 0.5+0.0im 0.0+0.0im 0.0+0.0im;
        0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0-0.5im 0.0+0.0im 0.0+0.0im -0.5+0.0im -0.5+0.0im 0.0+0.5im 0.0+0.0im;
         0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0-0.5im 0.0+0.0im 0.0+0.0im 0.0+0.5im -0.5+0.0im 0.0+0.5im;
         0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0-0.5im 0.0+0.0im 0.0+0.0im 0.0+0.5im -0.5+0.0im]

    result_matrix42 = [0.495478+1.2483e-17im -0.162938+0.0285877im 0.0104949-0.055023im -0.000950804+0.00066496im;
     -0.162938-0.0285877im 0.271419-1.3821e-17im -0.000400553+0.107483im 0.00474894+0.00823693im;
     0.0104949+0.055023im -0.000400553-0.107483im 0.155099-2.82874e-18im 0.00737685+0.0355666im;
    -0.000950804-0.00066496im 0.00474894-0.00823693im 0.00737685-0.0355666im 0.0780051-2.51094e-18im]
    result_matrix10 = [0.34111+5.57946e-18im -0.121394+0.0872953im -0.0693429+0.0173517im -0.0338042-0.0543639im;
     -0.121394-0.0872953im 0.248919-4.2406e-17im 0.0565208+0.0998314im -0.0158367+0.0694652im
     -0.0693429-0.0173517im 0.0565208-0.0998314im 0.223367-1.11563e-17im 0.0486222+0.0838723im;
     -0.0338042+0.0543639im -0.0158367-0.0694652im 0.0486222-0.0838723im 0.186604-2.90591e-18im]
    state = zeros(4,4)
    state[3,3] = 1
    #trivial evolutions
    @test evolve(zeros(16, 16), A, 0.) ≈ A
    @test evolve(zeros(16, 16), sparse(A), 0.) ≈ A
    @test evolve(spzeros(16, 16), A, 0.) ≈ A
    @test evolve(spzeros(16, 16), A, [0., 5., 10.])[1] ≈ A
    @test evolve(spzeros(16, 16), A, [0., 5., 10.])[2] ≈ A
    @test evolve(spzeros(16, 16), A, [0., 5., 10.])[3] ≈ A

    @test evolve(S, A, 0.) ≈ A
    @test evolve(Matrix{Float64}(I, 16, 16), A) ≈ A
    @test isapprox(evolve(S, A, 4.2), result_matrix42, atol=1e-5)
    @test isapprox(evolve(S, A, [0., 4.2, 10.]), [A, result_matrix42, result_matrix10], atol=1e-5)

  end
end


# 2020.04.01, Антонюк В.А.

include("HeuristicAlgorithm.jl")

function UpperSum(A::Array{Int,2})
    S,(H,W) = 0, size(A)
    @assert W == H
    for i=1:H, j=i+1:W  S += A[i,j] end
    return S
end

ReOrder(M::Matrix{Int}, V::Vector{Int}) = [M[i,j] for i ∈ V, j ∈ V]

function RecursiveIteration!(ScoreMatrix::Matrix{Int}, ToDrop=Vector{Int}())
    (H,W) = size(ScoreMatrix)
    @assert W == H
    TrueDropped = copy(ToDrop)
    TrueIndices!(TrueDropped)
    if length(TrueDropped) < H
        Summed = zeros(Int, H-length(TrueDropped))
        Dropped = 0
        for i = 1:H
            i ∈ TrueDropped && continue
            Dropped += 1
            for j = 1:W
                j ∈ TrueDropped && continue
                Summed[Dropped] += ScoreMatrix[i,j]
            end
        end
        (Val,Idx) = findmin(Summed)
        while Idx != nothing
            push!(ToDrop,Idx)
            if length(ToDrop) < H
                RecursiveIteration!(ScoreMatrix, ToDrop)
            else
                Dropped = copy(ToDrop)
                TrueIndices!(Dropped)
                LoopCount = 0
                while true
                    LoopCount += 1
                    OldDropped = copy(Dropped)
                    print("$Dropped : ")
                    print(UpperSum(ReOrder(ScoreMatrix, Dropped)))
                    OptimizeRanking!(Dropped,ScoreMatrix)
                    LoopCount < 100 || break
                    Dropped != OldDropped || break
                    print(" -> ")
                end
                println("\n")
            end
            pop!(ToDrop)
            Idx = findnext(x->x==Val, Summed, Idx+1)
        end
    end
end


#
# Metropolis-Hastings for logistic regression
# with missing input values
#

using PyPlot, PyCall
using Distributions
PyDict(matplotlib["rcParams"])["mathtext.fontset"] = "cm"
PyDict(matplotlib["rcParams"])["mathtext.rm"] = "serif"
PyDict(matplotlib["rcParams"])["lines.linewidth"] = 1.5
PyDict(matplotlib["rcParams"])["font.family"] = "TakaoPGothic"

function visualize_contour(samples_W, X, Y, text)
    R = 100
    xmin = minimum(X[1,:])
    xmax = maximum(X[1,:])
    ymin = minimum(X[2,:])
    ymax = maximum(X[2,:])
    lx = xmax - xmin
    ly = ymax - ymin
    xmin = xmin - 0.25 * lx
    xmax = xmax + 0.25 * lx
    ymin = ymin - 0.25 * ly
    ymax = ymax + 0.25 * ly

    x1 = linspace(xmin,xmax,R)
    x2 = linspace(ymin,ymax,R)
    x1grid = repmat(x1, 1, R)
    x2grid = repmat(x2', R, 1)
    val = [x1grid[:] x2grid[:]]'

    z_list = []
    for n in 1 : size(samples_W, 2)
        W = samples_W[:,n]
        z_tmp = [sigmoid(W'*val[:,i]) for i in 1 : size(val, 2)]
        push!(z_list, z_tmp)
    end
    z = mean(z_list)

    zgrid = reshape(z, R, R)

    # precition
    contour(x1grid, x2grid, zgrid, alpha=0.5, cmap=get_cmap("bwr"))
    scatter(X[1,Y.==1], X[2,Y.==1], c="r")
    scatter(X[1,Y.==0], X[2,Y.==0], c="b")
    xlim([xmin, xmax])
    ylim([ymin, ymax])
    title("$text")
end

function draw_line(W, xmin, xmax)
    y1 = - xmin*W[1]/W[2]
    y2 = - xmax*W[1]/W[2]
    plot([xmin, xmax], [y1, y2], c="k")
end

function sigmoid(x)
    return 1.0 / (1.0 + exp(-x[1]))
end

function bern_sample(mu)
    i = rand(Bernoulli(mu))
    val = zeros(2)
    val[i+1] = 1
    return val
end

function insert_X(X, X_miss)
    X_copy = deepcopy(X)
    X_copy[isnan(X)] = X_miss
    return X_copy
end

function MH(Y, X, M, Sigma_w, sigma_x, sigma_mh, max_iter, burnin)
    function ln_p_tilde(W, Y, X, Sigma_w, sigma_x)
        lkh = [Y[n]*log(sigmoid(W'*X[:, n])) + (1 - Y[n])*log(1 - sigmoid(W'*X[:, n])) for n in 1 : size(X, 2)]
        pri_w = -0.5 * (W'*inv(Sigma_w)*W)[1]
        pri_x = -0.5 * (sum(X.^2) * inv(sigma_x^2)) # Warning: Observation is not needed.
        return sum(lkh) + pri_w + pri_x
    end

    M = size(X, 1)
    N_miss = sum(isnan(X))

    ## sampling
    W = randn(M)
    X_miss = randn(N_miss)

    samples_W = Array{Float64, 2}(M, max_iter)
    samples_X = Array{Float64, 2}(N_miss, max_iter)
    for i in 1 : max_iter
        # sample W
        W_new = rand(MvNormal(W, sigma_mh^2*eye(M)))
        X_tmp = insert_X(X, X_miss)
        r = exp(minimum([ln_p_tilde(W_new, Y, X_tmp, Sigma_w, sigma_x) - ln_p_tilde(W, Y, X_tmp, Sigma_w, sigma_x), 0]))
        W = (rand() < r) ? W_new : W
        samples_W[:, i] = W

        # sample X_miss
        X_miss_new = rand(MvNormal(X_miss, sigma_mh^2*eye(N_miss)))
        X_tmp1 = insert_X(X, X_miss_new)
        X_tmp2 =  insert_X(X, X_miss)
        r = exp(minimum([ln_p_tilde(W, Y, X_tmp1, Sigma_w, sigma_x) - ln_p_tilde(W, Y, X_tmp2, Sigma_w, sigma_x), 0]))
        X_miss = (rand() < r) ? X_miss_new : X_miss
        samples_X[:, i] = X_miss
    end

    return samples_W[:,burnin+1:end], samples_X[:,burnin+1:end]
end

function test()
    ## compare results by discarding, interpolation and full observatino

    ########################
    # create model

    # hyperparameters
    M = 2
    Sigma_w = 20.0 * eye(M)

    ########################
    # create toy-data

    # sample parameters
    W = rand(MvNormal(zeros(M), Sigma_w))
    sigma_x = 1.0

    # sample data
    N = 30
    X_raw = 2 * rand(M, N) - 1.0
    Y = [rand(Bernoulli(sigmoid(W'*X_raw[:, n]))) for n in 1 : N]

    # missing values
    rate = 0.50
    X = deepcopy(X_raw)
    for n in 1 : N
        for m in 1 : M
            X[m,n] = (rand() < rate) ? NaN : X[m,n]
        end
    end

    idx = vec(!any(isnan(X), 1))
    X_red = X[:,idx]
    Y_red = Y[idx]

    ########################
    # inference
    sigma_mh = 0.10
    max_iter = 4000
    burnin = 1000

    samples_W1, samples_X1 = MH(Y_red, X_red, M, Sigma_w, sigma_x, sigma_mh, max_iter, burnin)
    samples_W2, samples_X2 = MH(Y, X, M, Sigma_w, sigma_x, sigma_mh, max_iter, burnin)
    samples_W3, samples_X3 = MH(Y, X_raw, M, Sigma_w, sigma_x, sigma_mh, max_iter, burnin)

    ########################
    # summarize
    idx_nan = vec(any(isnan(X), 1))
    X_est = insert_X(X, mean(samples_X2, 2))
    X_std = 0*deepcopy(X)
    X_std = insert_X(X_std, std(samples_X2, 2))

    ########################
    # visualize
    figure("result", figsize=(12,3))
    clf()
    subplot(1,3,1)
    visualize_contour(samples_W1, X_red, Y_red, "Reduced data (N=$(length(Y_red)))")

    subplot(1,3,2)
    Y_tmp = Y[idx_nan]
    X_tmp = X_est[:, idx_nan]
    scatter(X_tmp[1,Y_tmp.==1], X_tmp[2,Y_tmp.==1], c="r", marker="x")
    scatter(X_tmp[1,Y_tmp.==0], X_tmp[2,Y_tmp.==0], c="b", marker="x")
    for i in find(idx_nan)
        if Y[i] == 1
            plot([X_est[1,i] + X_std[1,i], X_est[1,i] - X_std[1,i]], [X_est[2,i], X_est[2,i]], "r", alpha=0.20, linewidth=5)
            plot([X_est[1,i], X_est[1,i]], [X_est[2,i] + X_std[2,i], X_est[2,i] - X_std[2,i]], "r", alpha=0.20, linewidth=5)
        else
            plot([X_est[1,i] + X_std[1,i], X_est[1,i] - X_std[1,i]], [X_est[2,i], X_est[2,i]], "b", alpha=0.20, linewidth=5)
            plot([X_est[1,i], X_est[1,i]], [X_est[2,i] + X_std[2,i], X_est[2,i] - X_std[2,i]], "b", alpha=0.20, linewidth=5)
        end
    end
    visualize_contour(samples_W2, X, Y, "Incomplete data (N=$N)")

    subplot(1,3,3)
    visualize_contour(samples_W3, X_raw, Y, "Full data (N=$N)")

    figure("traces", figsize=(12, 12))
    clf()
    subplot(4,1,1)
    plot(samples_W1')
    subplot(4,1,2)
    plot(samples_W2')
    subplot(4,1,3)
    plot(samples_X2')
    subplot(4,1,4)
    plot(samples_W3')

end

test()


push!(LOAD_PATH, string(pwd(),"/src"))

import AnalyticDomains
using ModifiedStokesSolver
using Base.Test

@testset "Discretization" begin
    numpanels = 100
    panelorder = 16
    curve = AnalyticDomains.starfish(n_arms = 5, amplitude=0.3);
    dcurve = discretize(curve, numpanels, panelorder);
    # Test equal arclength discretization
    h = sum(reshape(dcurve.dS, 16, :), 1)
    err = norm(h-h[1], Inf)
    @show err
    @test err < 1e-13
end


using Revise
using ADSeismic
using ADCME
using JLD2
ADCME.options.customop.verbose = false
mpi_init()
M = N = Int(sqrt(mpi_size()))
n = 100

param = MPIElasticPropagatorParams(NX= n * M, NY= n * N, n=n, NSTEP=2000, DELTAT=1e-4/2,
                        DELTAX=1.0, DELTAY=1.0)
compute_PML_Params!(param)
vp = 3300.
vs = 3300. / 1.732
rho = 2800.
λ, μ, ρ = compute_lame_parameters(param, vp, vs, rho)
source = Ricker(param, 50.0, 200.0)

srci = [div(param.NX,5)]
srcj = [div(param.NY,2)]
srctype = [0]
srcv = reshape(source, :, 1)
src = MPIElasticSource(param, srci, srcj, srctype, srcv)

rcvj = Array(20:n*param.M-19)
rcvi = ones(Int64, length(rcvj))*((n*param.N)÷5)
rcvtype = zeros(Int64, length(rcvj))
receiver =  MPIElasticReceiver(param, rcvi, rcvj, rcvtype)

propagator = MPIElasticPropagatorSolver(param, src, ρ, λ, μ)
MPISimulatedObservation!(propagator, receiver)

sess = Session(); init(sess)
@info "starting..."
rcvv, Vx = run(sess, [receiver.rcvv, propagator.vx])

@info "saving..."
@save "data/dat$(mpi_rank())-$(mpi_size()).jld2" rcvv

if mpi_size()>1
    mpi_finalize()
end
# if mpi_rank()==0 && mpi_size()==1
#     close("all")
#     pcolormesh(rcvv, cmap="gray")
#     colorbar()
#     xlabel("Location")
#     ylabel("Time")
#     savefig("receiver.png")

#     close("all")
#     visualize_wavefield(Vx, param)

# end


using Test
using Stripeline
using Dates
using Healpix
const Sl = Stripeline

eps = 3e-4 # Corresponds to pointing precision of 1.08 arcseconds

t_start = DateTime(2018, 01, 01, 0, 0, 0)

crab_az_stellarium_rad = 3.168070267969909
crab_alt_stellarium_rad = 1.4612458881566615

dirs = [π / 2 - crab_alt_stellarium_rad 2π - crab_az_stellarium_rad]

crab_ra_astropy_rad = 1.4596726619436968
crab_dec_astropy_rad = 0.3842255081802917
crab_position = sqrt(crab_ra_astropy_rad^2 + crab_dec_astropy_rad^2)

# Invert Crab coordinates into telescope pointing directions
groundq = telescopetoground(_->(0, deg2rad(20), 0), 0)
rotmatr = rotationmatrix_normalized(groundq)
vector = Healpix.ang2vec(dirs...)
dir = inv(rotmatr) * vector

# Compute skydirs
(skydirs, skyψ) = genpointings(_->(0, deg2rad(20), 0),
                               dir,
                               [0],
                               t_start,
                               latitude_deg = TENERIFE_LATITUDE_DEG,
                               longitude_deg = TENERIFE_LONGITUDE_DEG,
                               height_m = TENERIFE_HEIGHT_M)
crab_position_skydirs = sqrt(skydirs[1]^2 + skydirs[2]^2)
@test skydirs[1] ≈ crab_dec_astropy_rad atol = eps
@test skydirs[2] ≈ crab_ra_astropy_rad atol = eps
@test crab_position_skydirs ≈ crab_position atol = eps

t_1 = DateTime(2019, 02, 01, 2, 0, 0)
t_2 = DateTime(2019, 02, 02, 2, 0, 0)
t_3 = DateTime(2019, 02, 03, 2, 0, 0)

days = [t_1, t_2, t_3]

dirs = [π / 2 - 0.6174703397894518 2π - 4.852881197778698 ;
        π / 2 - 0.6024799007695448 2π - 4.859519751514131 ;
        π / 2 - 0.5875049757874335 2π - 4.866139397516001]

crab_ra_astropy_rad = 1.4596726619436968
crab_dec_astropy_rad = 0.3842255081802917
crab_position = sqrt(crab_ra_astropy_rad^2 + crab_dec_astropy_rad^2)

# Invert crab coordinates into telescope pointing directions
skydirs = Array{Float64}(undef, 3, 2)
for (idx, day) in enumerate(days)
    vector = Healpix.ang2vec(dirs[idx, 1], dirs[idx, 2])
    dir = inv(rotmatr) * vector

    (skydirections, skyψ) = genpointings(_->(0, deg2rad(20), 0),
                                         dir,
                                         [0],
                                         day,
                                         latitude_deg = TENERIFE_LATITUDE_DEG,
                                         longitude_deg = TENERIFE_LONGITUDE_DEG,
                                         height_m = TENERIFE_HEIGHT_M)

    skydirs[idx, 1] = skydirections[1]
    skydirs[idx, 2] = skydirections[2]
end

crab_position_skydirs = sqrt.(skydirs[:, 1].^2 + skydirs[:, 2].^2)

@test skydirs[1, 1] ≈ crab_dec_astropy_rad atol = eps
@test skydirs[1, 2] ≈ crab_ra_astropy_rad atol = eps
@test skydirs[2, 1] ≈ crab_dec_astropy_rad atol = eps
@test skydirs[2, 2] ≈ crab_ra_astropy_rad atol = eps
@test skydirs[3, 1] ≈ crab_dec_astropy_rad atol = eps
@test skydirs[3, 2] ≈ crab_ra_astropy_rad atol = eps
@test crab_position_skydirs[1] ≈ crab_position atol = eps
@test crab_position_skydirs[2] ≈ crab_position atol = eps
@test crab_position_skydirs[3] ≈ crab_position atol = eps

# Test flag `ground`
db = Sl.InstrumentDB()

time_duration = 86400
sampling_rate = 50.0
spin_velocity = 1

τ_s = 1 / sampling_rate
times = 0:τ_s:time_duration

(dirs, ψ) = genpointings(db.focalplane["I0"].orientation,
                         times;
                         latitude_deg = TENERIFE_LATITUDE_DEG) do time_s
    (0,
                              deg2rad(20.0),
                              Sl.timetorotang(time_s, spin_velocity))
end

expected_nsamples = convert(Int, time_duration * sampling_rate + 1)
@test size(dirs) == (expected_nsamples, 2)
@test size(ψ) == (expected_nsamples,)

(dirs, ψ) = genpointings(db.focalplane["I0"].orientation,
                         times;
                         ground = true,
                         latitude_deg = TENERIFE_LATITUDE_DEG) do time_s
    (0,
                              deg2rad(20.0),
                              Sl.timetorotang(time_s, spin_velocity))
end

expected_nsamples = convert(Int, time_duration * sampling_rate + 1)
@test size(dirs) == (expected_nsamples, 4)
@test size(ψ) == (expected_nsamples, 2)

# Check that all the values in the 3rd column are the same
@test dirs[1:end - 1, 3] ≈ dirs[2:end, 3]

# We strieve for an accuracy of a few arcseconds
@test rad2deg(dirs[1, 3]) ≈ 20.0 atol = 1e-3

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

# Check that the sign of the wheel angles is correct

let defaultdb = InstrumentDB()
    wheelfn(time_s) = (0, deg2rad(20), Stripeline.timetorotang(time_s, 1))

    # We do not start from t = 0s because we want the longitude to be positive
    timerange = 7200:1:(7200 + 60)
    G0_vec = defaultdb.focalplane["G0"].orientation
    V0_vec = defaultdb.focalplane["V0"].orientation

    dirG0, psiG0 = Stripeline.genpointings(wheelfn, G0_vec, timerange)
    dirV0, psiV0 = Stripeline.genpointings(wheelfn, V0_vec, timerange)

    # We expect G0 to draw larger circles in the sky
    @test minimum(dirG0[:, 1]) < minimum(dirV0[:, 1])
    @test minimum(dirG0[:, 1]) < minimum(dirV0[:, 1])
    @test maximum(dirG0[:, 1]) > maximum(dirV0[:, 1])
    @test maximum(dirG0[:, 2]) > maximum(dirV0[:, 2])
end


using SolidState
using LinearAlgebra
using StaticArrays
using Test

# types.jl

@test vec3d(1, 0, 0) ⋅ vec3d(0, 1, 0) ≈ 0
@test vec3d(2, 1, 3) ⋅ vec3d(2, 1, 3) ≈ 14
@test vec3d(1, 0, 0) × vec3d(0, 1, 0) ≈ vec3d(0, 0, 1)

@test norm2(vec3d(1, 0, 0)) ≈ 1
@test norm2(vec3d(1, 2, 3)) ≈ 14

# transformations.jl

@test rotate(π/6, vec3d(1, 0, 0)) ≈ rotatex(π/6)
@test rotate(π/6, vec3d(0, 1, 0)) ≈ rotatey(π/6)
@test rotate(π/6, vec3d(0, 0, 1)) ≈ rotatez(π/6)

@test translate(0.1, 0.2, 0.3) ≈ translate(vec3d(0.1, 0.2, 0.3))
@test translate(1, 0, 0) * translate(0, 2, 0) * translate(0, 0, 3) ≈ translate(1, 2, 3)

@test scalex(2) * scaley(3) * scalez(4) ≈ scale(2, 3, 4)

# Vectors should not vary when translated
@test translate(4, 5, 6) * vec3d(1, 2, 3) ≈ vec3d(1, 2, 3)
@test rotatez(π/2) * vec3d(1, 0, 0) ≈ vec3d(0, 1, 0)

# Remember that a ray transforms under the *inverse* transformation, so
# vector [1, 0, 0] rotates by -π/2 around the y axis
ray = rotatey(π/2) * Ray(vec3d(0, 0, 0), vec3d(1, 0, 0))
@test ray.origin ≈ vec3d(0, 0, 0)
@test ray_at_dist(ray, 10) ≈ vec3d(0, 0, 10)
@test ray(10) ≈ vec3d(0, 0, 10)

# cameras.jl

(u, v) = xy_to_uv(127, 348, 1024, 768)
(x, y) = uv_to_xy(u, v, 1024, 768)
@test [x, y] ≈ [127.0, 348.0]

camera = PerspectiveCamera(vec3d(0, 0, -2), # Position
                           vec3d(0, 0, 0),  # Look at
                           vec3d(0, 1, 0),  # Up
                           vec3d(1, 0, 0))  # Right

x, y, visible = project_point(camera, vec3d(0, 0, 0), 1024, 768)
@test x ≈ 1024 / 2
@test y ≈ 768 / 2
@test visible

x, y, visible = project_point(camera, vec3d(0, 0, -2), 1024, 768)
@test !visible

camera = OrthoCamera(vec3d(0, 0, -2), # Position
                     vec3d(0, 1, 0),  # Up
                     vec3d(1, 0, 0))  # Right

x, y, visible = project_point(camera, vec3d(0, 0, 0), 1024, 768)
@test x ≈ 1024 / 2
@test y ≈ 768 / 2
@test visible


####################################
# Load libraries
using DelimitedFiles
using Plots
using DataFrames
using CSV
using LaTeXStrings

####################################
# Load functions
include("bils_chang_kim_2011/steady_state_equilibrium.jl")

####################################
# Approximate value, policy, and wage functions,
# reservation productivity function, and stationary distributions
# given θ = 1
para            = ModelParams()
W, U, J, emp_policy, unemp_policy, wage = SolveWage(para)
x_star          = ReservationProductivity(para, W, U)
π_emp, π_unemp  = FindStationaryMeasures(para, emp_policy, unemp_policy, x_star)

####################################
# Save all objects externally
writedlm( "bils_chang_kim_2011/data/W.csv",  W, ',')
writedlm( "bils_chang_kim_2011/data/U.csv",  U, ',')
writedlm( "bils_chang_kim_2011/data/J.csv",  J, ',')
writedlm( "bils_chang_kim_2011/data/emp_policy.csv",  emp_policy, ',')
writedlm( "bils_chang_kim_2011/data/unemp_policy.csv",  unemp_policy, ',')
writedlm( "bils_chang_kim_2011/data/wage.csv",  wage, ',')
writedlm( "bils_chang_kim_2011/data/x_star.csv",  x_star, ',')
writedlm( "bils_chang_kim_2011/data/pi_emp.csv",  π_emp, ',')
writedlm( "bils_chang_kim_2011/data/pi_unemp.csv",  π_unemp, ',')

####################################
# Import objects (if necessary)
## Value of employment
W = CSV.read("bils_chang_kim_2011/data/W.csv", DataFrame, header = false)
W = Matrix(W)
## Value of unemployment
U = CSV.read("bils_chang_kim_2011/data/U.csv", DataFrame, header = false)
U = Matrix(U)
## Value of a matched job
J = CSV.read("bils_chang_kim_2011/data/J.csv", DataFrame, header = false)
J = Matrix(J)
## Asset policy function under employment
emp_policy      = CSV.read("bils_chang_kim_2011/data/emp_policy.csv", DataFrame, header = false)
emp_policy      = Matrix(emp_policy)
## Asset policy function under unemployment
unemp_policy    = CSV.read("bils_chang_kim_2011/data/unemp_policy.csv", DataFrame, header = false)
unemp_policy    = Matrix(unemp_policy)
## Wage function
wage            = CSV.read("bils_chang_kim_2011/data/wage.csv", DataFrame, header = false)
wage            = Matrix(wage)
## Employed distribution
pi_emp          = CSV.read("bils_chang_kim_2011/data/pi_emp.csv", DataFrame, header = false)
pi_emp          = Matrix(pi_emp)
## Unemployed distribution
pi_unemp        = CSV.read("bils_chang_kim_2011/data/pi_unemp.csv", DataFrame, header = false)
pi_unemp        = Matrix(pi_unemp)

####################################
# Plot all objects and save externally
## Save asset grid for x-axes
agrid = para.agrid
xgrid = para.xgrid
N_a = length(agrid)
N_x = length(xgrid)
## Plot value of employment for each x
fig1 = plot(title = "Asset Policy Under Employment");
for i in 1:9
    plot!(fig1, agrid, emp_policy[:,i], label="x=$(round(xgrid[i],digits=2))");
end
xlabel!("Current Assets");
ylabel!("Future Assets");
## Plot value of unemployment for each x
fig2 = plot(title = "Asset Policy Under Unemployment");
plot!(fig2, agrid, unemp_policy, label="");
xlabel!("Current Assets");
ylabel!("Future Assets");
## Plot value of firm-worker match for each x
fig3 = plot(title = "Matched Firm Value");
for i in 1:9
    plot!(fig3, agrid, J[:,i], label="x=$(round(xgrid[i],digits=2))");
end
xlabel!("Level of Asset Holdings");
ylabel!(L"J");
## Plot wages for each x
fig4 = plot(title = "Wage Schedule");
for i in 1:9
    plot!(fig4, agrid, wage[:,i], label="x=$(round(xgrid[i],digits=2))");
end
xlabel!("Level of Asset Holdings");
ylabel!("Wage");
## Plot W - U for each x
fig5 = plot(title = "Employment vs. Unemployment");
for i in 1:9
    plot!(fig5, agrid, W[:,i] - U, label="x=$(round(xgrid[i],digits=2))");
end
xlabel!("Level of Asset Holdings");
ylabel!(L"W-U");
## Plot asset distributions for employed and unemployed
fig6 = plot(title = "Asset Distributions for Employed and Unemployed");
emp_distr   = zeros(length(agrid))
unemp_distr = zeros(length(agrid))
for i in 1:N_a
    emp_distr[i]    = sum([π_emp[i+N_a*(j-1)] for j in 1:N_x])
    unemp_distr[i]  = sum([π_unemp[i+N_a*(j-1)] for j in 1:N_x])
end
plot!(fig6, agrid, emp_distr, label = "Employed");
plot!(fig6, agrid, unemp_distr, label = "Unemployed");
xlabel!("Level of Asset Holdings");
## Plot reservation match productivity
fig7 = plot(title = "Worker Assets and Reservation Match Productivity");
plot!(fig7, agrid, x_star, label = L"x^\ast(a)");
xlabel!("Level of Asset Holdings");
ylabel!("Productivity");
## All plots together
plot(fig1)
plot(fig2)
plot(fig3)
plot(fig4)
plot(fig5)
plot(fig6)
plot(fig7)
## Export plots
savefig(fig1, "bils_chang_kim_2011/figures/fig1.png")
savefig(fig2, "bils_chang_kim_2011/figures/fig2.png")
savefig(fig3, "bils_chang_kim_2011/figures/fig3.png")
savefig(fig4, "bils_chang_kim_2011/figures/fig4.png")
savefig(fig5, "bils_chang_kim_2011/figures/fig5.png")
savefig(fig6, "bils_chang_kim_2011/figures/fig6.png")
savefig(fig7, "bils_chang_kim_2011/figures/fig7.png")


## PARTE 1
cd(@__DIR__)
conv(c) = c == '.' ? 0 : c == '#' ? 1 : missing
#init = conv.(reduce(hcat, collect.(eachline("input.txt"))))
init = conv.([
    [ '.' '#' '.' ];
    [ '.' '.' '#' ];
    [ '#' '#' '#' ]
])
sz = 3 * max(size(init)...)
cells = zeros(Int, sz, sz, sz)
cells[1:size(init, 1),1:size(init, 2),1] .= init;
new_cells = similar(cells);
start = deepcopy(cells);

##
wrap(x, y, z; sz=size(cells)) = CartesianIndex((mod(i, 1:si) for (i,si) in zip((x, y, z), sz))...)
function neighbours(c; sz=size(cells))
    x, y, z = c.I .- 1
    Δs = (-1, 0, 1)
    [wrap(x+a+1, y+b+1, z+c+1; sz) for a in Δs, b in Δs, c in Δs if (a,b,c) != (0,0,0)] |> vec
end
isactive(c) = c == 1
function conway!(new_cells, cells)
    @assert size(new_cells) == size(cells)
    _isactive(c) = isactive(cells[c])
    foreach(CartesianIndices(cells)) do c
        state = cells[c]
        ns = neighbours(c; sz=size(cells))
        num = count(_isactive, ns)
        cells[c] = if _isactive(c)
            num != 2 && num != 3 ? 0 : state
        else
            num == 3 ? 1 : state
        end
    end
end

##
cells = deepcopy(start)
new_cells .= 0
println("Epoch (0): active cells $(count(isactive, cells))")
foreach(1:6) do epoch
    global cells
    global new_cells
    conway!(new_cells, cells)
    new_cells, cells = cells, new_cells
    println("Epoch ($epoch): active cells $(count(isactive, cells))")
end


using LinearAlgebra
using jInv.Mesh
using Test

# setup 3D Mesh
nc = [5, 7, 8]
x0 = [0, 0, 0]
domain = [x0[1], 1, x0[2], 1, x0[3], 1]
domain2D = [x0[1], 1, x0[2], 1]
h   = (domain[2:2:end]-domain[1:2:end])./nc
h1  = h[1]*ones(nc[1])
h2  = h[2]*ones(nc[2])
h3  = h[3]*ones(nc[3])

Mt = getTensorMesh3D(h1,h2,h3,x0)
Mr = getRegularMesh(domain,nc)

Mr2D = getRegularMesh(domain2D, [nc[1], nc[2]])

Meshes = (Mt, Mr, Mr2D);

for k=1:length(Meshes)

    M = Meshes[k]
    print("\ttesting boundary nodes for $(typeof(M))...")

    if M.dim==2

        X = getNodalGrid(M)[:,1];
        Y = getNodalGrid(M)[:,2];
        ib, iin = getBoundaryNodes(M);

        # Check if there are 0s in X and Y at boundary nodes
        @test length(X[ib]).-length(findall( x->(x < eps(Float32)), abs.(X[ib])))>0
        @test length(Y[ib]).-length(findall( x->(x < eps(Float32)), abs.(Y[ib])))>0

        # Check if there are 1s in X and Y at boundary nodes
        @test length(X[ib]).-length(findall( x->(x < eps(Float32)), abs.(X[ib].-1)))>0
        @test length(Y[ib]).-length(findall( x->(x < eps(Float32)), abs.(Y[ib].-1)))>0

        # Check that there aren't 0s in X, Y at inner nodes
        @test length(X[iin]).-length(findall(x->x!=0, X[iin]))==0
        @test length(Y[iin]).-length(findall(x->x!=0, Y[iin]))==0

        # Check that there aren't 1s in X, Y at inner nodes
        @test length(X[iin]).-length(findall(x->x!=1,X[iin]))==0
        @test length(Y[iin]).-length(findall(x->x!=1,Y[iin]))==0

    elseif M.dim==3

        X = getNodalGrid(M)[:,1];
        Y = getNodalGrid(M)[:,2];
        Z = getNodalGrid(M)[:,3];
        ib, iin = getBoundaryNodes(M);

        # Check if there are 0s in X, Y, and Z at boundary nodes
        @test length(X[ib])-length(findall( x->(x < eps(Float32)), abs.(X[ib])))>0
        @test length(Y[ib])-length(findall( x->(x < eps(Float32)), abs.(Y[ib])))>0
        @test length(Y[ib])-length(findall( x->(x < eps(Float32)), abs.(Z[ib])))>0

        # Check if there are 1s in X, Y, and Z at boundary nodes
        @test length(X[ib])-length(findall( x->(x < eps(Float32)), abs.(X[ib].-1)))>0
        @test length(Y[ib])-length(findall( x->(x < eps(Float32)), abs.(Y[ib].-1)))>0
        @test length(Y[ib])-length(findall( x->(x < eps(Float32)), abs.(Z[ib].-1)))>0

        # Check that there aren't 0s in X,Y,Z at inner nodes
        @test length(X[iin]).-length(findall(x->x!=0,X[iin]))==0
        @test length(Y[iin]).-length(findall(x->x!=0,Y[iin]))==0
        @test length(Z[iin]).-length(findall(x->x!=0,Z[iin]))==0

        # Check that there aren't 1s in X,Y,Z at inner nodes
        @test length(X[iin]).-length(findall(x->x!=1, X[iin]))==0
        @test length(Y[iin]).-length(findall(x->x!=1, Y[iin]))==0
        @test length(Z[iin]).-length(findall(x->x!=1, Z[iin]))==0

    end

    println("passed!\n")
end


using Poly
using Modeler

## Build problem objects
V1 = Variable("VOLT_1", Complex)
V2 = Variable("VOLT_2", Complex)

obj= 192 * V1*conj(V1) - (96.2+481im) * conj(V2)*V1 + (-96.2+481im) * conj(V1)*V2

Load2 = (-96.2 -481im) * conj(V1)*V2 + ( 96.2 +481im) * conj(V2)*V2
Load2_lb = -350 + 350im; Load2_ub = -350 + 350im
Unit1 = ( 96.2 +481im) * conj(V1)*V1 + (-96.2 -481im) * conj(V2)*V1
Unit1_lb = -400im; Unit1_ub = 600 + 400im
Voltm1 = abs2(V1)
Voltm1_lb = 0.90; Voltm1_ub = 1.10
Voltm2 = abs2(V2)
Voltm2_lb = 0.90; Voltm2_ub = 1.05

## Assemble problem
pb = Problem()
add_variable!(pb, V1)
add_variable!(pb, V2)

set_objective!(pb, obj)

add_constraint!(pb, "UNIT_1", Unit1_lb << Unit1 << Unit1_ub)
add_constraint!(pb, "LOAD_2", Load2_lb << Load2 << Load2_ub)
add_constraint!(pb, "VOLTM_1", Voltm1_lb << Voltm1 << Voltm1_ub)
add_constraint!(pb, "VOLTM_2", Voltm2_lb << Voltm2 << Voltm2_ub)

println(pb)

## Applying at a local solution
pt = Point([V1, V2], [-0.635304 + 0.706321im, 0.380961 + 0.898656im])

println("⟶ Constraints at Knitro local OPF solution:")
for (cstrName, cstr) in get_constraints(pb)
  println(cstrName, ": ", cstr.lb, " ≤ ", evaluate(cstr.p, pt), " ≤ ", cstr.ub)
end

println("⟶ Objective value: ", evaluate(pb.objective, pt))
println("⟶ Corresponding slacks: ")
print(get_slacks(pb, pt))
println("... and min slack: ", get_minslack(pb, pt))

## In real variables:
pb_R = pb_cplx2real(pb)
pt_R = Point([Variable("VOLT_1_r", Real),  Variable("VOLT_1_i", Real),  Variable("VOLT_2_r", Real),  Variable("VOLT_2_i", Real)], [-0.635304, 0.706321, 0.380961, 0.898656])

println("\n⟶ Constraints at Knitro local OPF solution, real problem:")
for (cstrName, cstr) in get_constraints(pb_R)
  println(cstrName, ": ", cstr.lb, " ≤ ", evaluate(cstr.p, pt_R), " ≤ ", cstr.ub)
end

println("⟶ Objective value, real problem: ", evaluate(pb_R.objective, pt_R))


"""
    pplus, eplus, pminus, eminus = dssplit2d(u, ABCD, p)

Split a convex polygon `p` into "plus" and "minus" polygons.
C(`pplus`) + D₁`u` ≥ 0, C(`pminus`) + D₁`u` ≤ 0.

`p` is given as a sequential list of vertices with the first vertex replicated
at the end of the list. `ABCD` describes the modulator structure, and `u` is the
modulator input.
"""
function dssplit2d(u, ABCD, p)
    n = size(ABCD, 1) - 1
    C = ABCD[n+1, 1:n]
    D1 = ABCD[n+1, n+1]	# D2=ABCD(n+1,n+2) must be zero
    N = size(p, 2)
    D1u = D1 * u
    y = p'*C + D1u*ones(N)

    sign1 = sign(y[1])
    i = findall(x -> sign(x) != sign1, y)
    i1 = i[1]       # First change of sign.
    pa = dscut(p[:,i1-1], y[i1-1], p[:,i1], y[i1])
    i2 = i[length(i)]   # Second change of sign.
    pb = dscut(p[:,i2], y[i2], p[:,i2+1], y[i2+1])
    if sign1 > 0
        pminus = [pa p[:,i] pb pa]
        pplus = [p[:,1:i1-1] pa pb p[:,i2+1:N]]
    else
        pplus = [pa p[:,i] pb pa ]
        pminus = [p[:,1:i1-1] pa pb p[:,i2+1:N]]
    end
    ne = size(pplus, 2)
    eplus = [(1:ne)'; (2:ne)' 1]
    ne = size(pminus, 2)
    eminus = [(1:ne)'; (2:ne)' 1]

    return pplus, eplus, pminus, eminus
end


mutable struct HexGrid

    node_size::Float64
    grid::Array{Array{Float64}}

    function HexGrid(res::Float64; dilat::Float64 = 1.0, theta::Float64 = 0.0,
                     bor_x::Array{Float64,1} = [0.0,1.0], bor_y::Array{Float64,1} = [0.0,1.0],
                     node_size::Float64 = 1.0, scale::Bool = true)
        grid = ref_grid(res)
        if dilat != 1.0
            dilat_grid!(grid, dilat)
        end
        if theta != 0
            rotate_grid!(grid, theta)
        end
        if scale
            scale_grid!(grid,bor_x[1], bor_x[2], bor_y[1], bor_y[2])
        end
        new(node_size,grid)
    end

end

function ref_grid(hex_size::Float64)
    w = sqrt(3)*hex_size
    h = 2*hex_size
    hi,lo = 2,-1
    center_x = [i*w/2+lo for i in 0:div(hi-lo, w/2)]
    center_y = [j*0.75*h+lo for j in 0:div(hi-lo, 0.75*h)]
    grid = Array{Float64}[]
    for i in eachindex(center_x)
        for j in eachindex(center_y)
            if (i+j)%2==0; push!(grid, [center_x[i], center_y[j]]); end
        end
    end
    return grid
end

function rotate_grid!(hex_grid::Array{Array{Float64}}, theta::Float64)
    rot = [cos(theta) -sin(theta); sin(theta) cos(theta)]
    center = [0.5,0.5]
    for i in eachindex(hex_grid)
        hex_grid[i] .= rot*(hex_grid[i] .- center) .+ center
    end
    nothing
end

function dilat_grid!(hex_grid::Array{Array{Float64}},ratio::Float64)
    for i in eachindex(hex_grid)
        hex_grid[i] .*= ratio
    end
    nothing
end

function dilat_grid!(hex_grid::HexGrid,ratio::Float64)
    for i in eachindex(hex_grid.grid)
        hex_grid.grid[i] .*= ratio
    end
    hex_grid.node_size *= ratio
    nothing
end

function trunc_grid!(hex_grid::Array{Array{Float64}}, lo_x::Float64 = 0.0,
                     hi_x::Float64 = 1.0, lo_y::Float64 = 0.0,
                     hi_y::Float64 = 1.0)
    filter!(x->(lo_x<=x[1]<=hi_x) && (lo_y<=x[2]<=hi_y),hex_grid)
    nothing
end

function scale_grid!(hex_grid::Array{Array{Float64}}, lo_x::Float64, hi_x::Float64,
                     lo_y::Float64, hi_y::Float64)
    l1 = hi_x - lo_x
    l2 = hi_y - lo_y
    if l1 >= l2
        len = l1
        trunc_grid!(hex_grid,0.0,1.0,0.0,l2/l1)
    end
    if l2 > l1
        len = l2
        trunc_grid!(hex_grid,0.0,l1/l2)
    end
    lo = [lo_x,lo_y]
    for i in eachindex(hex_grid)
        hex_grid[i] .= hex_grid[i] .*len .+lo
    end
    nothing
end

function scale_grid!(hex_grid::HexGrid, lo_x::Float64, hi_x::Float64,
                     lo_y::Float64, hi_y::Float64)
    l1 = hi_x - lo_x
    l2 = hi_y - lo_y
    if l1 >= l2
        len = l1
        trunc_grid!(hex_grid.grid,0.0,1.0,0.0,l2/l1)
    end
    if l2 > l1
        len = l2
        trunc_grid!(hex_grid.grid,0.0,l1/l2)
    end
    lo = [lo_x,lo_y]
    for i in eachindex(hex_grid.grid)
        hex_grid.grid[i] .= hex_grid.grid[i] .*len .+lo
    end
    hex_grid.node_size *= len
    nothing
end

function dist(p1::Array{Float64},p2::Array{Float64})
    d = sqrt((p1[1]-p2[1])^2+(p1[2]-p2[2])^2)
    return d
end

# %% hexmodule struct
mutable struct HexModule

    n_grid::Int64
    grids::Array{HexGrid}

    function HexModule(res::Float64; dilat::Float64 = 1.0, theta::Float64 = 0.0,
                       bor_x::Array{Float64,1} = [0.0,1.0], bor_y::Array{Float64,1} = [0.0,1.0],
                       mod_size::Int64 = 2) # parameter
        ref = HexGrid(res, scale = false)
        grids = HexGrid[]
        for i in 1:mod_size^2
            new_grid = deepcopy(ref)
            trans_grid!(new_grid, i, mod_size)
            dilat_grid!(new_grid, dilat)
            rotate_grid!(new_grid.grid, theta)
            scale_grid!(new_grid,bor_x[1], bor_x[2], bor_y[1], bor_y[2])
            push!(grids, new_grid)
        end
        new(mod_size^2,grids)
    end
end

function trans_grid!(hex_grid::HexGrid, placement::Int64, mod_size::Int64)
    D = dist(hex_grid.grid[1],hex_grid.grid[2])
    d = D/sqrt(3)
    d_trans = d/(mod_size)
    theta_i = pi/3
    theta_j = -pi/3
    i = div(placement-1,mod_size)
    j = (placement-1)%mod_size
    translation_i = d_trans*[cos(theta_i), sin(theta_i)]
    translation_j = d_trans*[cos(theta_j), sin(theta_j)]
    for _ in 1:i
        for index in eachindex(hex_grid.grid)
            hex_grid.grid[index] .+= translation_i
        end
    end
    for _ in 1:j
        for index in eachindex(hex_grid.grid)
            hex_grid.grid[index] .+= translation_j
        end
    end
    hex_grid.node_size = d_trans/2
    nothing
end

# %% Visualization
using Plots
gr()

function scatter_grid(hex_grid::HexGrid)
    x = []
    y = []
    for point in hex_grid.grid
        push!(x,point[1])
        push!(y,point[2])
    end
    display(scatter(x,y,color = :red))
    nothing
end

function scatter_grid!(hex_grid::HexGrid)
    x = []
    y = []
    for point in hex_grid.grid
        push!(x,point[1])
        push!(y,point[2])
    end
    display(scatter!(x,y,legend = false, xlabel = "Variable 1", ylabel = "Variable 2"))
    nothing
end

function scatter_module(hex_mod::HexModule)
    for i in 1:hex_mod.n_grid
        scatter_grid!(hexmod.grids[i])
    end
end

# %% hexmodule test

borx = [0.0,1.0]
bory = [0.0,1.0]
hexmod = HexModule(0.03, dilat = 1.0, theta = pi/10, bor_x = borx,
                   bor_y = bory, mod_size = 2)

plot(size = (600,600))
scatter_grid!(hexmod.grids[1])
#
#
# # %%
#
# @show hexmod.grids[1].grid[1]
# @show hexmod.grids[1].grid[2]
# @show hexmod.grids[1].grid[3]
# @show hexmod.grids[1].grid[4]


using BilevelBenchmark
using Test

@testset "Bilevel unconstrained" begin
    D_ul = 5
    D_ll = 5
    for i = 1:8
        x, y = SMD_solutions(D_ul, D_ll, i)
        Fx, fy = SMD_leader(x, y, i), SMD_leader(x, y, i)
        @test  abs(Fx) + abs(fy) < 1e-10
    end
end

@testset "Bilevel Ranges" begin
    D_ul = 5
    D_ll = 5
    for i = 1:8
        bounds_ul, bounds_ll = SMD_ranges(D_ul, D_ll, i)
        x = bounds_ul[1,:] + (bounds_ul[2,:] - bounds_ul[1,:]) .* rand(D_ul)
        y = bounds_ll[1,:] + (bounds_ll[2,:] - bounds_ll[1,:]) .* rand(D_ul)

        Fx, fy = SMD_leader(x, y, i), SMD_leader(x, y, i)
        @test  abs(Fx) + abs(fy) > -Inf
    end
end

@testset "SMD" begin
    p, q, r, s = 3, 3, 2, 0
    for i = 1:8
        @test SMD_leader(ones(1000), ones(1000), i, p, q, r, s) ≈ SMD_leader(ones(1000), ones(1000), i, p, q, r, s)
        @test SMD_follower(ones(1000), ones(1000), i, p, q, r, s) ≈ SMD_follower(ones(1000), ones(1000), i, p, q, r, s)
    end
end

@testset "TP" begin
    for i = 1:10
        @test TP_test(i) >= 0

        # bounds_ul, bounds_ll = TP_ranges(i)

    end
end

@testset "PMM" begin
    for i = 1:6
        F, f = BilevelBenchmark.PMM_test(10, 10, i)
        @test abs(f + F) <= 1e-16
    end

    x = zeros(9)
    for fnum = 1:6
        y = PMM_Ψ(x, 9, fnum)
        if fnum < 7
            F = PMM_leader(x, y, fnum)
            f = PMM_follower(x, y, fnum)
        else
            F, _ = PMM_leader(x, y, fnum)
            f, _ = PMM_follower(x, y, fnum)
        end

        @test abs(F) <= 1e-16 &&  abs(f) <= 1e-16
    end
end

@testset "bilevel constrained" begin
    D_ul = 5
    D_ll = 5

    for i = 9:12
        p,q,r,s,lenG, leng = SMD_settings(D_ul, D_ll, i)

        bounds_ul, bounds_ll = SMD_ranges(D_ul, D_ll, i)
        x = bounds_ul[1,:] + (bounds_ul[2,:] - bounds_ul[1,:]) .* rand(D_ul)
        y = bounds_ll[1,:] + (bounds_ll[2,:] - bounds_ll[1,:]) .* rand(D_ul)

        F, G = SMD_leader(x, y, i)
        f, g = SMD_follower(x, y, i)

        @test length(G) == lenG
        @test length(g) == leng
        @test abs(F + sum(G)) >= 0
        @test abs(f + sum(g)) >= 0
    end
end

@testset "bilevel get problem" begin
    for benchmark in [:SMD, :PMM, :TP]
        for fnum = 1:6
            F, f, bounds_ul, bounds_ll, Ψ, z = get_problem(benchmark, fnum)
            x = bounds_ul[1,:]
            y = bounds_ll[1,:]
            Fxy = F(x, y)
            fxy = f(x,y)
            @test abs(z[1]) + abs(z[2]) >= 0
        end
    end

end


using GeometricIntegrators
using GeometricExamples
using ImportMacros

@import GeometricProblems.LotkaVolterra2d as problem

include("lotka_volterra_2d_settings_firk.jl")

tableau_list = get_tableau_list_firk()

equ = problem.lotka_volterra_2d_ode()

include("lotka_volterra_2d.jl")

run_tableau_list(tableau_list, equ, Δt, ntime)

generate_html(tableau_list, figures, figures_detail, prefix, fig_suff, fig_dir, out_dir, html_file;
              title=coll_title, params=get_params(Δt, ntime, problem))


using SymPy, Plots
gr()

#TODO consider looking into how Michal Kaveslov did path following -> discrete points (states), but how did he know which ones to choose??

# track data
type Track
  R
  x0
  y0
  Xn
  Yn
end

# EX4 test data -> https://www.varsitytutors.com/hotmath/hotmath_help/topics/shortest-distance-between-a-point-and-a-circle
function Track()
  Track(7,
        4,
       -5,
       -4,
       -11);
 end

function initTrack(;R=R,x0=x0,y0=y0,Xn=Xn,Yn=Yn)
  t=Track();
  t.R=R;
  t.x0=x0;
  t.y0=y0;
  t.Xn=Xn;
  t.Yn=Yn;
  return t
end

function minD(t::Track)
  sign(sqrt((t.Xn-t.x0)^2 + (t.Yn-t.y0)^2)-t.R)*(sqrt((t.Xn-t.x0)^2 + (t.Yn-t.y0)^2)-t.R);
end

function plotTrack(t::Track)
  pp=plot(0,leg=:false);
  title!(@sprintf("min D (m) = %0.2f",minD(t)))
  pts = Plots.partialcircle(0,2π,100,t.R);
  x, y = Plots.unzip(pts)
  x += t.x0;
  y += t.y0;
  pts = collect(zip(x, y));
  plot!(pts,line=(:path,5,:solid,:green),label="Track");
  XX=[t.Xn;t.x0];YY=[t.Yn;t.y0];
  plot!(XX,YY,line=(:path,5,:solid),label="D # ");

  pts = Plots.partialcircle(0,2π,100,0.25);
  x, y = Plots.unzip(pts);
  x += t.Xn;  y += t.Yn;
  pts = collect(zip(x, y));
  plot!(pts,c=:black,fill=:true,label="Vehicle")
  #xlabel!("X (m)");  ylabel!("Y (m)");
  #xlims!(-10,10); ylims!(-15,10); #TODO --> use this as a position plot
end

format=:png;

# ex4
t4=Track();
plotTrack(Track(t4));
savefig(string("t4.",format));

# ex3
t3=initTrack(;R=5,x0=-3,y0=3,Xn=-2,Yn=0);
plotTrack(t3);
savefig(string("t3.",format));

# real track
t1=initTrack(;R=100,x0=100,y0=100,Xn=0,Yn=0);
plotTrack(t1);
savefig(string("t1.",format));

# this script is designed to calcualte the analytical distance from a point to an elipse (if there is one...)
# not able to find the analytical solution

function elllipseD()
  a,b,x,y=symbols("a,b,x,y",positive=true);
  x0,y0,Xn,Yn=symbols("x0,y0,Xn,Yn",positive=true);
  D=symbols("D",positive=true);
  epsilon=symbols("epsilon",positive=true)
  # equation for distance
  exp1=(x-Xn)^2 + (y-Yn)^2 - D^2;

  # derivative of D w.r.t x == slope of D
  exp3=diff(exp1,x) - (y-Yn)/(x-Xn+sign(x-Xn)*epsilon);  # only Chuck Noris can divide by zero

  # equation for ellipse
  exp2=(((x-x0)/a)^2 + ((y-y0)/b)^2)-1;

  # both equations
  exp = [exp1, exp2, exp3];

  # solve the system
  solve(exp,[D,x])

  ####################
  # test the solution
  ####################
  # test data - > all variables and parameters must be positive
  A_=5; B_=10;
  x0_=5; y0_=20;
  Xn_=10; Yn_=0;

  # find the nearest point
  solve(exp1,)# solve for x

  pp=plot(0,leg=:false);

  A=5; B=10;
  x0=5; y0=20;
  Xn=10; Yn=0;

  pts = Plots.partialcircle(0,2π,100,A)
  x, y = Plots.unzip(pts)
  x += x0;
  y = B/A*y + y0;
  #pts = collect(zip(x, y))
  plot!(Shape(x,y), c=:red,leg=true,label="Obstacle",leg=:bottomleft)
  scatter!((Xn,Yn), markershape = :square, markercolor = :black, markersize = 4,fill = (0, 1, :black),leg=true, grid=true,label="Vehicle")
end


struct JC69abs <: JC69
    λ::Float64
    function JC69abs(λ::Float64)
        if λ <= 0.
            error("JC69 parameter λ must be positive")
        end
        new(λ)
    end
end

function show(io::IO, object::JC69abs)
  print(io, "\r\e[0m\e[1mJ\e[0mukes and \e[1mC\e[0mantor 19\e[1m69\e[0m model (absolute rate form)
λ = $(object.λ)")
end

JC69(λ) = JC69abs(λ)

function Q(mod::JC69abs)
  λ = mod.λ
  Q₁ =  0.25 * λ
  Q₂ = -(Q₁ * 3)
  return Qmatrix(Q₂, Q₁, Q₁, Q₁,
                 Q₁, Q₂, Q₁, Q₁,
                 Q₁, Q₁, Q₂, Q₁,
                 Q₁, Q₁, Q₁, Q₂)
end

function P(mod::JC69abs, t::Float64)
  if t < 0.0
    error("t must be positive")
  end
  λ = mod.λ
  ω = exp(-t * λ)
  P₁ = 0.25 + 0.75 * ω
  P₂ = 0.25 + 0.25 * ω
  return Pmatrix(P₁, P₂, P₂, P₂,
                 P₂, P₁, P₂, P₂,
                 P₂, P₂, P₁, P₂,
                 P₂, P₂, P₂, P₁)
end


# OptiGas: Optimize Gas Network using Benders Algorithm
# Version: 2.0
# Author: Edward J. Xu, edxu96@outlook.com
# Date: March 20th, 2019
push!(LOAD_PATH, "$(homedir())/Desktop/OptiGas, DTU42136")
cd("$(homedir())/Desktop/OptiGas, DTU42136")
# using BendersMilp_EDXU
using Benders
using LinearAlgebra
# -------------------------------------------------------------------
# 1. Parameters
vec_nameNodes = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J"]
numNodes = length(vec_nameNodes)
vec_xNode = [1    3    8    2    4    5    6    7    9    9]
vec_yNode = [1    2    1    7    9    5    3    7    9    4]
mat_arcTwoNodes = [0    1    0    0    0    0    0    0    0    0;
                   0    0    0    0    0    1    0    0    0    0;
                   0    0    0    0    0    0    0    0    0    1;
                   1    0    0    0    1    0    0    0    0    0;
                   0    0    0    0    0    1    0    0    0    0;
                   1    0    0    1    0    0    0    1    0    0;
                   0    1    1    0    0    0    0    0    0    0;
                   0    0    0    0    0    0    0    0    1    0;
                   0    0    0    0    1    0    0    0    0    0;
                   0    0    0    0    0    1    1    0    1    0]
vec_netEject =  [0    0   0    0    0   90    0    0   0   80 ]
vec_netInject = [10   30  20   10   40   0   25   20  15    0 ]
mat_distance = zeros(Float64, numNodes, numNodes)
for m = 1: numNodes
    for n = 1: numNodes
        mat_distance[m, n] = floor(sqrt( (vec_xNode[m] - vec_xNode[n]) * (vec_xNode[m] - vec_xNode[n]) +
                                   (vec_yNode[m] - vec_yNode[n]) * (vec_yNode[m] - vec_yNode[n]) ) )
    end
end
mat_fixedCost = zeros(Float64, numNodes, numNodes)
for m = 1: numNodes
    for n = 1: numNodes
        mat_fixedCost[m, n] = 10 * mat_distance[m, n]
    end
end
# ----------------------------------------------------------------------------------------------------------------------
# Question 4
# Fomulation of Matrix A
mat_a1_1 = zeros(numNodes^2, numNodes^2)
for mn = 1: numNodes^2
    mat_a1_1[mn, mn] = -1
end
# mat_a2_1 = zeros(numNodes^2, numNodes^2)
mat_a3_1 = zeros(numNodes, numNodes^2)
seq_zeroTenNinety = collect(0:10:90)
for m = 1: numNodes
    mat_a3_1[m, (10 * m - 9): (10 * m)] = repeat([-1], 10)  # sent out
    mat_a3_1[m, (seq_zeroTenNinety + repeat([m],10))] = repeat([1], 10)  # sent in
    mat_a3_1[:, (11 * m - 10)] = repeat([0], 10)  # self-sending doesn't count
end
mat_a_1 = vcat(mat_a1_1, mat_a3_1)
# Fomulation of Matrix B
mat_b1_1 = zeros(numNodes^2, numNodes^2)
for mn = 1: numNodes^2
    mat_b1_1[mn, mn] = 170
end
# mat_b2_1 = Diagnal(repeat([1], numNodes^2))
mat_b3_1 = zeros(numNodes, numNodes^2)
mat_b_1 = vcat(mat_b1_1, mat_b3_1)
# Fomulation of Vector b
vec_b1_1 = zeros(numNodes^2)
# vec_b2_1 = zeros(numNodes^2)
# for m = 1: numNodes
#     for n = 1: numNodes
#         vec_b2_1[10 * (m - 1) + n] = mat_arcTwoNodes[m, n]
#     end
# end
vec_b3_1 = zeros(numNodes)
for n = 1: numNodes
    vec_b3_1[n] = vec_netInject[n] - vec_netEject[n]
end
vec_b_1 = vcat(vec_b1_1, vec_b3_1)
vec_b_1 = hcat(vec_b_1)
#
vec_max_y_1 = repeat([1], numNodes^2)
vec_max_y_1 = hcat(vec_max_y_1)
vec_min_y_1 = zeros(numNodes^2)
for m = 1: numNodes
    for n = 1: numNodes
        vec_min_y_1[10 * (m - 1) + n] = mat_arcTwoNodes[m, n]
    end
end
#
vec_c_1 = zeros(numNodes^2)
for m = 1: numNodes
    for n = 1: numNodes
        vec_c_1[10 * (m - 1) + n] = mat_distance[m, n]
    end
end
vec_c_1 = hcat(vec_c_1)
#
vec_f_1 = zeros(numNodes^2)
for m = 1: numNodes
    for n = 1: numNodes
        vec_f_1[10 * (m - 1) + n] = mat_fixedCost[m, n]
    end
end
vec_f_1 = hcat(vec_f_1)
#
Benders.milp(n_x = numNodes^2,
             n_y = numNodes^2,
             vec_min_y = vec_min_y_1,
             vec_max_y = vec_max_y_1,
             vec_c = vec_c_1,
             vec_f = vec_f_1,
             vec_b = vec_b_1,
             mat_a = mat_a_1,
             mat_b = mat_b_1,
             epsilon = 0.0001,
             timesIterationMax = 1000)
# 1280 - sum(mat_arcTwoNodes .* mat_fixedCost)
# ----------------------------------------------------------------------------------------------------------------------
# Question 5
# Fomulation of Matrix A
mat_a1_2 = zeros(numNodes^2, numNodes^2)
for mn = 1: numNodes^2
    mat_a1_2[mn, mn] = -1
end
mat_a3_2 = zeros(numNodes, numNodes^2)
seq_zeroTenNinety = collect(0:10:90)
for m = 1: numNodes
    mat_a3_2[m, (10 * m - 9): (10 * m)] = repeat([-1], 10)  # sent out
    mat_a3_2[m, (seq_zeroTenNinety + repeat([m],10))] = repeat([1], 10)  # sent in
    mat_a3_2[:, (11 * m - 10)] = repeat([0], 10)  # self-sending doesn't count
end
mat_a_2 = vcat(mat_a1_2, mat_a3_2)
# Fomulation of Matrix B
mat_b1_2 = zeros(numNodes^2, numNodes^2)
for mn = 1: numNodes^2
    mat_b1_2[mn, mn] = 170
end
mat_b3_2 = zeros(numNodes, numNodes^2)
mat_b_2 = vcat(mat_b1_2, mat_b3_2)
# Fomulation of Vector b
vec_b1_2 = zeros(numNodes^2)
vec_b3_2 = zeros(numNodes)
for n = 1: numNodes
    vec_b3_2[n] = vec_netInject[n] - vec_netEject[n]
end
vec_b_2 = vcat(vec_b1_2, vec_b3_2)
vec_b_2 = hcat(vec_b_2)
#
vec_max_y_2 = repeat([1], numNodes^2)
vec_max_y_2 = hcat(vec_max_y_2)
#
vec_c_2 = zeros(numNodes^2)
for m = 1: numNodes
    for n = 1: numNodes
        vec_c_2[10 * (m - 1) + n] = mat_distance[m, n]
    end
end
vec_c_2 = hcat(vec_c_2)
#
vec_f_2 = zeros(numNodes^2)
for m = 1: numNodes
    for n = 1: numNodes
        vec_f_2[10 * (m - 1) + n] = mat_fixedCost[m, n]
    end
end
vec_f_2 = hcat(vec_f_2)
#
Benders.milp(n_x = numNodes^2,
            n_y = numNodes^2,
            vec_min_y = repeat([0], numNodes^2),
            vec_max_y = vec_max_y_2,
            vec_c = vec_c_2,
            vec_f = vec_f_2,
            vec_b = vec_b_2,
            mat_a = mat_a_2,
            mat_b = mat_b_2,
            epsilon = 0.0001,
            timesIterationMax = 1000)
# obj - sum(mat_arcTwoNodes .* mat_fixedCost)
# ----------------------------------------------------------------------------------------------------------------------
# using LightGraphs
# using GraphPlot
# using Compose
# using Cairo
# using Fontconfig
# mat_result_y = [0 1 0 0 0 0 0 0 0 0;
#                 0 0 0 0 0 1 0 0 0 0;
#                 0 0 0 0 0 0 0 0 0 1;
#                 1 0 0 0 1 0 0 0 0 0;
#                 0 0 0 0 0 1 0 0 0 0;
#                 1 1 0 1 1 0 0 1 0 0;
#                 0 1 1 0 0 0 0 0 0 0;
#                 0 0 0 0 0 0 0 0 1 0;
#                 0 0 0 0 1 0 0 0 0 0;
#                 0 0 1 0 0 1 1 1 1 0]
# graph_initial = DiGraph(mat_arcTwoNodes)
# plot_graph_initial = gplot(graph_initial, nodelabel = vec_nameNodes)
# graph_result = graph_initial
# mat_arcTwoNodes_new = mat_result_y - mat_arcTwoNodes
# for m = 1: numNodes
#     for n = 1: numNodes
#         if mat_arcTwoNodes_new[m,n] == 1
#             add_edge!(graph_result, m, n)
#         end
#     end
# end
# gplot(graph_result, nodelabel = vec_nameNodes)
# draw(PNG("graph_initial.png", 16cm, 16cm), gplot(graph_initial))
# draw(PNG("graph_result.png", 16cm, 16cm), gplot(result))

# Final Corrected
vec_x= [0 0 0 0 0 0 0 0 0 0
        0 0 0 0 0 0 0 0 0 0
        0 0 0 0 0 0 0 0 0 0
        0 0 0 0 0 0 0 0 0 0
        0 0 0 0 0 0 0 0 0 0
        10 10 0 10 40 0 0 20 0 0
        0 20 20 0 0 0 0 0 0 0
        0 0 0 0 0 0 0 0 0 0
        0 0 0 0 0 0 0 0 0 0
        0 0 0 0 0 0 65 0 15 0]
vec_y= [0 1 0 0 0 0 0 0 0 0
        0 0 0 0 0 1 0 0 0 0
        0 0 0 0 0 0 0 0 0 1
        1 0 0 0 1 0 0 0 0 0
        0 0 0 0 0 1 0 0 0 0
        1 1 0 1 1 0 0 1 0 0
        0 1 1 0 0 0 0 0 0 0
        0 0 0 0 0 0 0 0 1 0
        0 0 0 0 1 0 0 0 0 0
        0 0 0 0 0 1 1 0 1 0]

vec_y: [0 1 0 0 0 0 0 0 0 0
        0 0 0 0 0 1 0 0 0 0
        0 0 0 0 0 0 0 0 0 1
        1 0 0 0 1 0 0 0 0 0
        0 0 0 0 0 1 0 0 0 0
        1 1 0 1 1 0 0 1 0 0
        0 1 1 0 0 0 0 0 0 0
        0 0 0 0 0 0 0 0 1 0
        0 0 0 0 1 0 0 0 0 0
        0 0 1 0 0 1 1 1 1 0]
result_q: 600.0

vec_x= [0 0 0 0 0 0 0 0 0 0
        0 0 0 0 0 0 0 0 0 0
        0 0 0 0 0 0 0 0 0 0
        0 0 0 0 0 0 0 0 0 0
        0 0 0 0 0 0 0 0 0 0
        10 30 0 10 40 0 0 0 0 0
        0 0 0 0 0 0 0 0 0 0
        0 0 0 0 0 0 0 0 15 0
        0 0 0 0 0 0 0 0 0 0
        0 0 20 0 0 0 25 35 0 0]

# vec_result_x = reshape(transpose(vec_x), numNodes^2)
vec_result_node = zeros(numNodes)
for m = 1: numNodes
    vec_result_node[m] = sum(vec_x[(10 * m - 9): (10 * m)]) - sum(vec_x[(seq_zeroTenNinety + repeat([m],10))])
end
vec_y= [0 1 0 0 0 0 0 0 0 0
        0 0 0 0 0 1 0 0 0 0
        0 0 0 0 0 0 0 0 0 1
        1 0 0 0 1 0 0 0 0 0
        0 0 0 0 0 1 0 0 0 0
        1 1 0 1 1 0 0 1 0 0
        0 1 1 0 0 0 0 0 0 0
        0 0 0 0 0 0 0 0 1 0
        0 0 0 0 1 0 0 0 0 0
        0 0 1 0 0 1 1 1 1 0]


function count_bits(chars::AbstractVector{Char})
    c0 = count(chars .== '0')
    c1 = count(chars .== '1')
    c0, c1
end

function most_common_bit(chars::AbstractVector{Char})
    c0, c1 = count_bits(chars)
    c0 > c1 ? '0' : '1'
end

function least_common_bit(chars::AbstractVector{Char})
    c0, c1 = count_bits(chars)
    c0 <= c1 ? '0' : '1'
end

function oxygen_generator(chars::AbstractVector{Char})
    chars .== most_common_bit(chars)
end

function co2_scrubber(chars::AbstractVector{Char})
    chars .== least_common_bit(chars)
end

function rating(matrix::AbstractMatrix{Char}, bit_criteria)
    for bit_position in 1:size(matrix, 2)
        index = bit_criteria(matrix[:, bit_position])
        matrix = matrix[index, :]
        if size(matrix, 1) == 1
            return matrix[1, :]
        end
    end
end

function binary_number(chars::AbstractVector{Char})
    s = String(chars)
    parse(Int, s, base = 2)
end

matrix = vcat((permutedims(Vector{Char}(line)) for line in readlines())...)
gamma = binary_number([most_common_bit(col) for col in eachcol(matrix)])
epsilon = binary_number([least_common_bit(col) for col in eachcol(matrix)])
oxygen = binary_number(rating(matrix, oxygen_generator))
co2 = binary_number(rating(matrix, co2_scrubber))
println(gamma * epsilon)
println(oxygen * co2)