### Some methods for velocity calculation for a system driven by a surface force
"""
Interface velocity in a viscous liquid where regions seperated by interface do have the same viscousity. If γ is passed the curvattureless algorithm is used to take into account force due to surface tension. Returns velocity field projected on vertex normals. Usefull for equilibrium calculations.
"""
function stokesvelocity(points,normals,faces,forcen,etaP,gammap)
    vareas = vertexareas(points,faces)
    velocityn = zeros(Float64,length(points))

    for xkey in 1:length(points)

        x = points[xkey]
        nx = normals[xkey]
        fx = forcen[xkey]

        s = 0
        for ykey in 1:length(points)
            if ykey==xkey
                continue
            end

            y = points[ykey]
            ny = normals[ykey]
            fy = forcen[ykey]

            ### I will need to check a missing 2
            s += vareas[ykey]*1 ./8/pi/etaP* dot(y-x,nx+ny)/norm(y-x)^3*(1-3*dot(y-x,nx)*dot(y-x,ny)/norm(y-x)^2) * gammap

            ### ????????
            s += vareas[ykey]*1 ./8/pi/etaP* ( dot(nx,ny)/norm(x-y) + dot(nx,x -y)*dot(ny,x-y)/norm(x-y)^3 )*(fy - fx)
        end
        velocityn[xkey] = s
    end
    return velocityn
end

function stokesvelocity(points,normals,faces,forcen,etaP)

    vareas = vertexareas(points,faces)
    velocityn = zeros(Float64,size(points,2))

    for xkey in 1:length(points)

        x = points[xkey]
        nx = normals[xkey]
        fx = forcen[xkey]

        s = 0
        for ykey in 1:length(points)
            if ykey==xkey
                continue
            end

            y = points[ykey]
            ny = normals[ykey]
            fy = forcen[ykey]

            s += vareas[ykey]*1 ./8/pi/etaP* ( dot(nx,ny)/norm(x-y) + dot(nx,x -y)*dot(ny,x-y)/norm(x-y)^3 )*(fy - fx)
        end

        velocityn[xkey] = s
    end

    return velocityn
end


module Day21

using Memoize

import ..data_dir # from parent module
input = read(joinpath(data_dir, "day21"), String)

export part1, part2

function parse_player(player_string::AbstractString)
    numbers = [match.match for match in eachmatch(r"[[:digit:]]+", player_string)]
    Player(parse.(Int64, numbers)...)
end

function parse_players(input::AbstractString)
    player_strings = split(input, "\n", keepempty = false)
    parse_player.(player_strings)
end

mutable struct Player
    number::Int64
    position::Int64
    score::Int64
end
Player(number, position) = Player(number, position, 0)
Player(number) = Player(number, 0, 0)

new_position(position::Int64, movement::Int64) = mod(position + movement - 1, 10) + 1

function move!(player::Player, movement::Int64)
    player.position = new_position(player.position, movement)
    player.score += player.position
end

mutable struct DeterministicDice
    sequence::Vector{Int64}
    rolls::Int64
end
DeterministicDice(size::Int64) = DeterministicDice(collect(1:size), 0)

function roll!(dice::DeterministicDice)
    if dice.rolls == length(dice.sequence)
        dice.rolls = 0
    end
    dice.rolls += 1
    return dice.sequence[dice.rolls]
end

function roll!(dice::DeterministicDice, times::Int64)
    sum(roll!(dice) for i = 1:times)
end

function part1(input = input)
    players = parse_players(input)
    sort!(players, by = p -> p.number)
    dice = DeterministicDice(100)
    n_rolls = 0
    while true
        for player in players
            move!(player, roll!(dice, 3))
            n_rolls += 3
            if player.score >= 1000
                losing_player = first(filter(p -> p.number != player.number, players))
                return losing_player.score * n_rolls
            end
        end
    end
end

# enumerating the possible outcomes of rolling 3d3, the frequencies of each total:
const dice_sums = [(3, 1), (4, 3), (5, 6), (6, 7), (7, 6), (8, 3), (9, 1)]

# adapted from code by https://github.com/AxlLind
# https://github.com/AxlLind/AdventOfCode2021/blob/main/src/bin/21.rs
@memoize function dirac_dice_game(
    this_player_position::Int64,
    other_player_position::Int64,
    this_player_score::Int64 = 0,
    other_player_score::Int64 = 0
)
    if other_player_score >= 21
        return 0, 1 # other player wins
    end

    score = (0, 0)
    for (roll, times) in dice_sums
        this_player_new_position = new_position(this_player_position, roll)
        this_player_new_score = this_player_score + this_player_new_position
        # Note that the players alternate below
        other_player_wins, this_player_wins = dirac_dice_game(
            other_player_position,
            this_player_new_position,
            other_player_score,
            this_player_new_score
        )
        score = score .+ times .* (this_player_wins, other_player_wins)
    end

    score
end

function part2(input = input)
    players = parse_players(input)
    sort!(players, by = p -> p.number)
    wins = dirac_dice_game(players[1].position, players[2].position)
    maximum(wins)
end

end #module


module Regression

using Reexport
using ArrayViews
@reexport using EmpiricalRisks

import Base.LinAlg: BlasReal, axpy!
import Base.LinAlg.LAPACK: gels!, gelsy!, gelsd!
import EmpiricalRisks: value, value!, value_and_grad!

export
	# linearreg
	llsq,
	ridgereg,

	# regproblems
	UnivariateRegression,
	MultivariateRegression,

	linearreg,
	logisticreg,
	mlogisticreg,

	# solve
	GD,
	AGD,
	BFGS,
	LBFGS,

	# proxsolve
	ProximalDescent,
	ProxGD,
	ProxAGD

# source files

include("common.jl")
include("linearreg.jl")
include("regproblems.jl")
include("optimbase.jl")
include("solve.jl")
include("proxsolve.jl")
include("print.jl")

const default_solver = BFGS()

end # module


using NeuralQuantum, Random, Plots

Random.seed!(Random.GLOBAL_RNG, 1234)

# Parameters of the transverse field ising model
dims = [5,5]
h = 3.0
J = 1.0

# Constructs the Hilbert space for N 1//2 spins.
N    = prod(dims)
hilb = HomogeneousSpin(N, 1//2)

coord(i,j) = i + (j-1)*dims[2]

# Builds the hamiltonian
H = LocalOperator(hilb)
for i=1:dims[1]
    for j=1:dims[2]
        p = coord(i,j)
        global H  -= h * sigmax(hilb, p)
        hop = J * sigmaz(hilb, p) * sigmaz(hilb, coord(mod(i, dims[1])+1,j)) +
                J * sigmaz(hilb, p) * sigmaz(hilb, coord(i, mod(j, dims[2])+1))
        global H += hop
    end
end

net  = RBM(Float32, N, 1, af_logcosh)
init_random_pars!(net, sigma=0.01)

sampl = MetropolisSampler(LocalRule(), 75, N, burn=100)
algo  = SR(ϵ=(0.1), algorithm=sr_qlp)

is = BatchedSampler(net, sampl, H, algo;
                    batch_sz=16)

optimizer = Optimisers.Descent(0.1)

Evalues = Float64[];
Eerr = Float64[];
for i=1:1000
    ldata, prec = sample!(is)
    ob = compute_observables(is)

    println("$i - $ldata")

    push!(Evalues, real(ldata.mean))
    push!(Eerr, ldata.error)
    grad = precondition!(prec, algo, i)
    Optimisers.update!(optimizer, net, grad)
end

plot(Evalues, yerr=Eerr)

exact = -80.13310152422413
hline!([exact, exact])


# Tracing charged particle in a static analytic dipole magnetic field.
# This example shows how to trace protons of a certain energy in a analytic
# Earth-like magnetic dipole field. Pyplot is used for interactively changing
# views in 3D.
#
# Hongyang Zhou, hyzhou@umich.edu

using TestParticle
using TestParticle: getB_dipole, getE_dipole, sph2cart, Rₑ
using OrdinaryDiffEq
using PyPlot

Ek = 5e7 # [eV]

m = TestParticle.mᵢ
q = TestParticle.qᵢ
c = TestParticle.c
Rₑ = TestParticle.Rₑ

# initial velocity, [m/s]
v₀ = sph2cart(c*sqrt(1-1/(1+Ek*q/(m*c^2))^2), 0.0, π/4)
# initial position, [m]
r₀ = sph2cart(2.5*Rₑ, 0.0, π/2)
stateinit = [r₀..., v₀...]
# obtain field
param = prepare(getE_dipole, getB_dipole)
tspan = (0.0, 2.0)

prob = ODEProblem(trace_analytic!, stateinit, tspan, param)

sol = solve(prob, Tsit5(); save_idxs=[1,2,3])

x = getindex.(sol.u,1) / Rₑ
y = getindex.(sol.u,2) / Rₑ
z = getindex.(sol.u,3) / Rₑ

## Visualization

using3D()
fig = plt.figure()
ax = fig.gca(projection="3d")

plot(x,y,z)

for ϕ in range(0, stop=2*π, length=10)
   plot(fieldline(ϕ)..., color="r", alpha=0.3)
end


# The plant node
# This file is part of the openBuildNet framework developed at EPFL.
# (C) 2016 LA-EPFL
# Author: Truong X. Nghiem (truong.nghiem@gmail.com, xuan.nghiem@epfl.ch)

# See: http://users.isy.liu.se/johanl/yalmip/pmwiki.php?n=Examples.StandardMPC
# The plant model:
# x(t+1) = A*x(t) + B*u(t)
# y(t) = x(t)

using openbuildnet

# Model data
const A = [2. -1.;1. 0.2]
const B = [1. 0.]'
const x0 = [3., 1.];  # Initial state

x = x0  # The plant's state
xs = Vector{typeof(x)}()  # To store the state trajectory
push!(xs, x)

# Create the node
node = OBNNode("plant", "simplempc")

# Create the ports
u = create_input(node, "u", Vector{Float64})
y = create_output(node, "y", Vector{Float64})

# Create Callbacks
on_init(node) do
  println("INIT")
end

on_term(node) do
  println("TERMINATE")
end

const MAINBLOCK = 0

on_block_output(node, MAINBLOCK) do
  set(y, x)
end

on_block_state(node, MAINBLOCK) do
  global x
  x = A*x + B*get(u)
  push!(xs, x)  # save the state
end

println("Ready to start simulation...")
status = runsafe(node, 20)
println("Final status: ", status)

# Plot the results
if status == 2
  using PyPlot
  xsmat = hcat(xs...)' # convert to matrix
  plot(xsmat[:,1], xsmat[:,2])
  show()
end


using LinearAlgebra
import AtomicLevels: angular_integral

@testset "jj2ℓsj" begin
    @testset "Transform matrix" begin
        R = jj2ℓsj(ros"1[s] 2[s-p] k[s-d]"...)
        # Since it is a rotation matrix, its inverse is simply the
        # transpose.
        @test norm(inv(Matrix(R)) - R') < 10eps()
    end

    @testset "Transform individual spin-orbitals" begin
        for o in [o"5s", o"5p", o"5d"]
            for so in spin_orbitals(o)
                ℓ = so.orb.ℓ
                mℓ,ms = so.m
                mj = mℓ + ms
                pure = abs(mj) == ℓ + half(1)
                linear_combination = jj2ℓsj(so)
                @test length(linear_combination) == (pure ? 1 : 2)
                @test norm(last.(linear_combination)) ≈ 1
                @test all(isequal(mj), map(o -> first(o.m), first.(linear_combination)))
            end
        end
    end

    @testset "Angular integrals" begin
        @test isone(angular_integral(SpinOrbital(o"ks", (0,half(1))), SpinOrbital(o"1s", (0,half(1)))))
        @test iszero(angular_integral(SpinOrbital(o"ks", (0,half(1))), SpinOrbital(o"2p", (0,half(1)))))

        @test isone(angular_integral(SpinOrbital(ro"ks", half(1)), SpinOrbital(ro"1s", half(1))))
        @test iszero(angular_integral(SpinOrbital(ro"ks", half(1)), SpinOrbital(ro"2p-", half(1))))

        @test isone(angular_integral(SpinOrbital(o"ks", (0,half(1))), SpinOrbital(ro"1s", half(1))))
        @test isone(angular_integral(SpinOrbital(ro"1s", half(1)), SpinOrbital(o"ks", (0,half(1)))))

        @test angular_integral(SpinOrbital(o"2p", (0,half(1))), SpinOrbital(ro"2p-", half(1))) == clebschgordan(1, 0, half(1), half(1), half(1), half(1))
    end

    @testset "#orbitals = $(length(ros))" for (ros,nb) in ((rsos"l[p]", 4), (rsos"k[s-d] l[d]", 18), (filter(o -> o.m[1]==1//2, rsos"l[d]"), 1))
        os, bs = jj2ℓsj(ros)
        @test length(os) == length(ros)
        for (ro,o) in zip(ros, os)
            @test sum(o.m) == ro.m[1]
            @test o.orb.n == ro.orb.n
            @test o.orb.ℓ == ro.orb.ℓ
        end
        @test length(bs) == nb
        for (i,j,b) in bs
            if size(b,1) == 1
                @test i == j
            else
                @test i ≠ j
            end
            @test b == b'
            @test b^2 ≈ I
        end
    end
end


using Oceananigans
using Oceananigans.Units
using GLMakie

grid = RectilinearGrid(size=(128, 128), x=(0, 128), z=(-64, 0), topology=(Periodic, Flat, Bounded))

b_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(1e-8))

@inline growth_func(x, y, z, t, p) = 1 / p.τ * exp(z / p.h)
growth = Forcing(growth_func, parameters=(τ=1hour, h=4.0))
sinking = AdvectiveForcing(WENO5(), w=-1)

model = NonhydrostaticModel(; grid,
                            tracers = (:b, :P),
                            buoyancy = BuoyancyTracer(),
                            boundary_conditions = (; b=b_bcs),
                            forcing = (; P = (growth, sinking)))

bᵢ(x, y, z) = 1e-5 * z + 1e-9 * rand()
set!(model, b=bᵢ)

Δz = grid.Δzᵃᵃᶜ
Δt = 0.1 * Δz # for a sinking velocity w=1
simulation = Simulation(model; Δt, stop_iteration = 1000)

fig = Figure()
ax = Axis(fig[1, 1])

P = model.tracers.P
hm = heatmap!(ax, interior(P, :, 1, :))

function update!(sim)
    hm.input_args[1][] = interior(P, :, 1, :)
    return nothing
end

simulation.callbacks[:plot] = Callback(update!, IterationInterval(100))

display(fig)

run!(simulation)


module Manifolds

import ManifoldsBase:
    @trait_function,
    _access_nested,
    _get_basis,
    _injectivity_radius,
    _inverse_retract,
    _inverse_retract!,
    _read,
    _retract,
    _retract!,
    _write,
    active_traits,
    allocate,
    allocate_coordinates,
    allocate_result,
    allocate_result_type,
    allocation_promotion_function,
    array_value,
    base_manifold,
    check_point,
    check_size,
    check_vector,
    copy,
    copyto!,
    default_inverse_retraction_method,
    default_retraction_method,
    default_vector_transport_method,
    decorated_manifold,
    distance,
    dual_basis,
    embed,
    embed!,
    exp,
    exp!,
    get_basis,
    get_basis_default,
    get_basis_diagonalizing,
    get_basis_orthogonal,
    get_basis_orthonormal,
    get_basis_vee,
    get_component,
    get_coordinates,
    get_coordinates!,
    get_coordinates_diagonalizing,
    get_coordinates_diagonalizing!,
    get_coordinates_orthogonal,
    get_coordinates_orthonormal,
    get_coordinates_orthogonal!,
    get_coordinates_orthonormal!,
    get_coordinates_vee!,
    get_embedding,
    get_iterator,
    get_vector,
    get_vector!,
    get_vector_diagonalizing,
    get_vector_diagonalizing!,
    get_vector_orthogonal,
    get_vector_orthonormal,
    get_vector_orthogonal!,
    get_vector_orthonormal!,
    get_vectors,
    gram_schmidt,
    hat,
    hat!,
    injectivity_radius,
    _injectivity_radius,
    injectivity_radius_exp,
    inner,
    isapprox,
    is_point,
    is_vector,
    inverse_retract,
    inverse_retract!,
    _inverse_retract,
    _inverse_retract!,
    inverse_retract_caley!,
    inverse_retract_embedded!,
    inverse_retract_nlsolve!,
    inverse_retract_pade!,
    inverse_retract_polar!,
    inverse_retract_project!,
    inverse_retract_qr!,
    inverse_retract_softmax!,
    log,
    log!,
    manifold_dimension,
    mid_point,
    mid_point!,
    norm,
    number_eltype,
    number_of_coordinates,
    parallel_transport_along,
    parallel_transport_along!,
    parallel_transport_direction,
    parallel_transport_direction!,
    parallel_transport_to,
    parallel_transport_to!,
    parent_trait,
    power_dimensions,
    project,
    project!,
    representation_size,
    retract,
    retract!,
    retract_caley!,
    retract_exp_ode!,
    retract_pade!,
    retract_polar!,
    retract_project!,
    retract_qr!,
    retract_softmax!,
    set_component!,
    vector_space_dimension,
    vector_transport_along, # just specified in Euclidean - the next 5 as well
    vector_transport_along_diff,
    vector_transport_along_project,
    vector_transport_along!,
    vector_transport_along_diff!,
    vector_transport_along_project!,
    vector_transport_direction,
    vector_transport_direction_diff,
    vector_transport_direction!,
    vector_transport_direction_diff!,
    vector_transport_to,
    vector_transport_to_diff,
    vector_transport_to_project,
    vector_transport_to!,
    vector_transport_to_diff!,
    vector_transport_to_project!, # some overwrite layer 2
    _vector_transport_direction,
    _vector_transport_direction!,
    _vector_transport_to,
    _vector_transport_to!,
    vee,
    vee!,
    zero_vector,
    zero_vector!,
    CotangentSpace,
    TangentSpace
import Base:
    copyto!,
    convert,
    foreach,
    identity,
    in,
    inv,
    isempty,
    length,
    ndims,
    showerror,
    size,
    transpose

using Base.Iterators: repeated
using Distributions
using Einsum: @einsum
using HybridArrays
using Kronecker
using Graphs
using LinearAlgebra
using ManifoldsBase:
    ℝ,
    ℂ,
    ℍ,
    AbstractBasis,
    AbstractDecoratorManifold,
    AbstractInverseRetractionMethod,
    AbstractManifold,
    AbstractManifoldPoint,
    AbstractNumbers,
    AbstractOrthogonalBasis,
    AbstractOrthonormalBasis,
    AbstractPowerManifold,
    AbstractPowerRepresentation,
    AbstractRetractionMethod,
    AbstractTrait,
    AbstractVectorTransportMethod,
    AbstractLinearVectorTransportMethod,
    ApproximateInverseRetraction,
    ApproximateRetraction,
    CachedBasis,
    CayleyRetraction,
    CayleyInverseRetraction,
    ComplexNumbers,
    ComponentManifoldError,
    CompositeManifoldError,
    CotangentSpaceType,
    CoTFVector,
    DefaultBasis,
    DefaultOrthogonalBasis,
    DefaultOrthonormalBasis,
    DefaultOrDiagonalizingBasis,
    DiagonalizingBasisData,
    DiagonalizingOrthonormalBasis,
    DifferentiatedRetractionVectorTransport,
    EmbeddedManifold,
    EmptyTrait,
    ExponentialRetraction,
    FVector,
    IsIsometricEmbeddedManifold,
    IsEmbeddedManifold,
    IsEmbeddedSubmanifold,
    IsExplicitDecorator,
    LogarithmicInverseRetraction,
    ManifoldsBase,
    NestedPowerRepresentation,
    NestedReplacingPowerRepresentation,
    TraitList,
    NLSolveInverseRetraction,
    ODEExponentialRetraction,
    OutOfInjectivityRadiusError,
    PadeRetraction,
    PadeInverseRetraction,
    ParallelTransport,
    PolarInverseRetraction,
    PolarRetraction,
    PoleLadderTransport,
    PowerManifold,
    PowerManifoldNested,
    PowerManifoldNestedReplacing,
    ProjectedOrthonormalBasis,
    ProjectionInverseRetraction,
    ProjectionRetraction,
    ProjectionTransport,
    QuaternionNumbers,
    QRInverseRetraction,
    QRRetraction,
    RealNumbers,
    ScaledVectorTransport,
    SchildsLadderTransport,
    SoftmaxRetraction,
    SoftmaxInverseRetraction,
    TangentSpaceType,
    TCoTSpaceType,
    TFVector,
    TVector,
    ValidationManifold,
    ValidationMPoint,
    ValidationTVector,
    VectorSpaceType,
    VeeOrthogonalBasis,
    @invoke_maker,
    _euclidean_basis_vector,
    combine_allocation_promotion_functions,
    default_inverse_retraction_method,
    geodesic,
    merge_traits,
    next_trait,
    number_system,
    real_dimension,
    rep_size_to_colons,
    shortest_geodesic,
    size_to_tuple,
    trait
using Markdown: @doc_str
using Random
using RecipesBase
using RecipesBase: @recipe, @series
using Colors: RGBA
using Requires
using SimpleWeightedGraphs: AbstractSimpleWeightedGraph, get_weight
using StaticArrays
using Statistics
using StatsBase
using StatsBase: AbstractWeights
using RecursiveArrayTools: ArrayPartition

include("utils.jl")

include("product_representations.jl")
include("differentiation/differentiation.jl")
include("differentiation/riemannian_diff.jl")
include("differentiation/embedded_diff.jl")

# Main Meta Manifolds
include("manifolds/ConnectionManifold.jl")
include("manifolds/MetricManifold.jl")
include("manifolds/VectorBundle.jl")
include("groups/group.jl")

# Features I: Which are extended on Meta Manifolds
include("distributions.jl")
include("projected_distribution.jl")
include("statistics.jl")

# Meta Manifolds II: Products
include("manifolds/ProductManifold.jl")

METAMANIFOLDS = [
    AbstractManifold,
    AbstractDecoratorManifold,
    AbstractPowerManifold,
    PowerManifoldNested,
    PowerManifoldNestedReplacing,
    ProductManifold,
    TangentSpaceAtPoint,
    ValidationManifold,
    VectorBundle,
]

# Features II: That require metas
include("atlases.jl")
include("cotangent_space.jl")

# Meta Manifolds II: Power Manifolds
include("manifolds/PowerManifold.jl")
include("manifolds/GraphManifold.jl")

#
# Manifolds
#
include("manifolds/Euclidean.jl")
include("manifolds/Lorentz.jl")

include("manifolds/CenteredMatrices.jl")
include("manifolds/CholeskySpace.jl")
include("manifolds/Circle.jl")
include("manifolds/Elliptope.jl")
include("manifolds/FixedRankMatrices.jl")
include("manifolds/GeneralizedGrassmann.jl")
include("manifolds/GeneralizedStiefel.jl")
include("manifolds/Grassmann.jl")
include("manifolds/Hyperbolic.jl")
include("manifolds/MultinomialDoublyStochastic.jl")
include("manifolds/MultinomialSymmetric.jl")
include("manifolds/ProbabilitySimplex.jl")
include("manifolds/PositiveNumbers.jl")
include("manifolds/ProjectiveSpace.jl")
include("manifolds/Rotations.jl")
include("manifolds/SkewHermitian.jl")
include("manifolds/Spectrahedron.jl")
include("manifolds/Stiefel.jl")
include("manifolds/StiefelEuclideanMetric.jl")
include("manifolds/StiefelCanonicalMetric.jl")
include("manifolds/Sphere.jl")
include("manifolds/SphereSymmetricMatrices.jl")
include("manifolds/Symmetric.jl")
include("manifolds/SymmetricPositiveDefinite.jl")
include("manifolds/SymmetricPositiveDefiniteLinearAffine.jl")
include("manifolds/SymmetricPositiveDefiniteLogCholesky.jl")
include("manifolds/SymmetricPositiveDefiniteLogEuclidean.jl")
include("manifolds/SymmetricPositiveSemidefiniteFixedRank.jl")
include("manifolds/Tucker.jl")
include("manifolds/Symplectic.jl")
include("manifolds/SymplecticStiefel.jl")

# Product or power based manifolds
include("manifolds/Torus.jl")
include("manifolds/Multinomial.jl")
include("manifolds/Oblique.jl")
include("manifolds/EssentialManifold.jl")

#
# Group Manifolds
include("groups/GroupManifold.jl")

# a) generics
include("groups/addition_operation.jl")
include("groups/multiplication_operation.jl")
include("groups/connections.jl")
include("groups/metric.jl")
include("groups/group_action.jl")
include("groups/group_operation_action.jl")
include("groups/validation_group.jl")
include("groups/product_group.jl")
include("groups/semidirect_product_group.jl")

# Special Group Manifolds
include("groups/general_linear.jl")
include("groups/special_linear.jl")
include("groups/translation_group.jl")
include("groups/special_orthogonal.jl")
include("groups/circle_group.jl")
include("groups/heisenberg.jl")

include("groups/translation_action.jl")
include("groups/rotation_action.jl")

include("groups/special_euclidean.jl")

@doc raw"""
    Base.in(p, M::AbstractManifold; kwargs...)
    p ∈ M

Check, whether a point `p` is a valid point (i.e. in) a [`AbstractManifold`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/types.html#ManifoldsBase.AbstractManifold)  `M`.
This method employs [`is_point`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/functions.html#ManifoldsBase.is_point) deactivating the error throwing option.
"""
Base.in(p, M::AbstractManifold; kwargs...) = is_point(M, p, false; kwargs...)

@doc raw"""
    Base.in(p, TpM::TangentSpaceAtPoint; kwargs...)
    X ∈ TangentSpaceAtPoint(M,p)

Check whether `X` is a tangent vector from (in) the tangent space $T_p\mathcal M$, i.e.
the [`TangentSpaceAtPoint`](@ref) at `p` on the [`AbstractManifold`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/types.html#ManifoldsBase.AbstractManifold)  `M`.
This method uses [`is_vector`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/functions.html#ManifoldsBase.is_vector) deactivating the error throw option.
"""
function Base.in(X, TpM::TangentSpaceAtPoint; kwargs...)
    return is_vector(base_manifold(TpM), TpM.point, X, false; kwargs...)
end

function __init__()
    @require FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" begin
        using .FiniteDifferences
        include("differentiation/finite_differences.jl")
    end

    @require OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" begin
        using .OrdinaryDiffEq: ODEProblem, AutoVern9, Rodas5, solve
        include("differentiation/ode.jl")
    end

    @require NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56" begin
        using .NLsolve: NLsolve
        include("nlsolve.jl")
    end

    @require Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" begin
        using .Test: Test
        include("tests/tests_general.jl")
        export test_manifold
        include("tests/tests_group.jl")
        export test_group, test_action
    end

    @require Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" begin
        using RecipesBase: @recipe, @series
        using Colors: RGBA
        include("recipes.jl")
    end

    @require RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01" begin
        @require Colors = "5ae59095-9a9b-59fe-a467-6f913c188581" begin
            using .RecipesBase: @recipe, @series
            using Colors: RGBA
            include("recipes.jl")
        end
    end

    return nothing
end

#
export CoTVector, AbstractManifold, AbstractManifoldPoint, TVector
export AbstractSphere, AbstractProjectiveSpace
export Euclidean,
    ArrayProjectiveSpace,
    ArraySphere,
    CenteredMatrices,
    CholeskySpace,
    Circle,
    Elliptope,
    EssentialManifold,
    FixedRankMatrices,
    GeneralizedGrassmann,
    GeneralizedStiefel,
    Grassmann,
    HeisenbergGroup,
    Hyperbolic,
    Lorentz,
    MultinomialDoubleStochastic,
    MultinomialMatrices,
    MultinomialSymmetric,
    Oblique,
    PositiveArrays,
    PositiveMatrices,
    PositiveNumbers,
    PositiveVectors,
    ProbabilitySimplex,
    ProjectiveSpace,
    Rotations,
    SkewHermitianMatrices,
    SkewSymmetricMatrices,
    Spectrahedron,
    Sphere,
    SphereSymmetricMatrices,
    Stiefel,
    SymmetricMatrices,
    SymmetricPositiveDefinite,
    SymmetricPositiveSemidefiniteFixedRank,
    Symplectic,
    SymplecticStiefel,
    SymplecticMatrix,
    Torus,
    Tucker
export HyperboloidPoint, PoincareBallPoint, PoincareHalfSpacePoint, SVDMPoint, TuckerPoint
export HyperboloidTVector,
    PoincareBallTVector, PoincareHalfSpaceTVector, UMVTVector, TuckerTVector
export AbstractNumbers, ℝ, ℂ, ℍ

# decorator manifolds
export AbstractDecoratorManifold
export IsIsometricEmbeddedManifold, IsEmbeddedManifold, IsEmbeddedSubmanifold
export IsDefaultMetric, IsDefaultConnection, IsMetricManifold, IsConnectionManifold
export ValidationManifold, ValidationMPoint, ValidationTVector, ValidationCoTVector
export CotangentBundle,
    CotangentSpaceAtPoint, CotangentBundleFibers, CotangentSpace, FVector
export AbstractPowerManifold,
    AbstractPowerRepresentation,
    ArrayPowerRepresentation,
    NestedPowerRepresentation,
    NestedReplacingPowerRepresentation,
    PowerManifold
export ProductManifold
export GraphManifold, GraphManifoldType, VertexManifold, EdgeManifold
export ProjectedPointDistribution, ProductRepr, TangentBundle, TangentBundleFibers
export TangentSpace, TangentSpaceAtPoint, VectorSpaceAtPoint, VectorSpaceType, VectorBundle
export VectorBundleFibers
export AbstractVectorTransportMethod,
    DifferentiatedRetractionVectorTransport, ParallelTransport, ProjectedPointDistribution
export PoleLadderTransport, SchildsLadderTransport
export ProductVectorTransport
export AbstractAffineConnection,
    AbstractConnectionManifold, ConnectionManifold, LeviCivitaConnection
export AbstractCartanSchoutenConnection,
    CartanSchoutenMinus, CartanSchoutenPlus, CartanSchoutenZero
export AbstractMetric,
    RiemannianMetric,
    LorentzMetric,
    EmbeddedManifold,
    EuclideanMetric,
    LinearAffineMetric,
    LogCholeskyMetric,
    LogEuclideanMetric,
    MinkowskiMetric,
    PowerMetric,
    ProductMetric,
    RealSymplecticMetric,
    ExtendedSymplecticMetric,
    CanonicalMetric,
    MetricManifold
export AbstractAtlas, RetractionAtlas
export AbstractVectorTransportMethod, ParallelTransport, ProjectionTransport
export AbstractRetractionMethod,
    CayleyRetraction,
    ExponentialRetraction,
    QRRetraction,
    PolarRetraction,
    ProjectionRetraction,
    SoftmaxRetraction,
    ODEExponentialRetraction,
    PadeRetraction,
    ProductRetraction,
    PowerRetraction
export AbstractInverseRetractionMethod,
    ApproximateInverseRetraction,
    ApproximateLogarithmicMap,
    CayleyInverseRetraction,
    LogarithmicInverseRetraction,
    QRInverseRetraction,
    PolarInverseRetraction,
    ProjectionInverseRetraction,
    SoftmaxInverseRetraction
export AbstractEstimationMethod,
    GradientDescentEstimation,
    CyclicProximalPointEstimation,
    GeodesicInterpolation,
    GeodesicInterpolationWithinRadius,
    ExtrinsicEstimation
export CachedBasis,
    DefaultBasis,
    DefaultOrthogonalBasis,
    DefaultOrthonormalBasis,
    DiagonalizingOrthonormalBasis,
    InducedBasis,
    ProjectedOrthonormalBasis
export ComponentManifoldError, CompositeManifoldError
export ×,
    allocate,
    allocate_result,
    base_manifold,
    bundle_projection,
    change_metric,
    change_metric!,
    change_representer,
    change_representer!,
    check_point,
    check_vector,
    christoffel_symbols_first,
    christoffel_symbols_second,
    christoffel_symbols_second_jacobian,
    convert,
    complex_dot,
    decorated_manifold,
    det_local_metric,
    distance,
    dual_basis,
    einstein_tensor,
    embed,
    embed!,
    exp,
    exp!,
    flat,
    flat!,
    gaussian_curvature,
    geodesic,
    get_default_atlas,
    get_component,
    get_embedding,
    grad_euclidean_to_manifold,
    grad_euclidean_to_manifold!,
    hat,
    hat!,
    identity_element,
    identity_element!,
    induced_basis,
    incident_log,
    injectivity_radius,
    inner,
    inverse_local_metric,
    inverse_retract,
    inverse_retract!,
    isapprox,
    is_default_connection,
    is_default_metric,
    is_group_manifold,
    is_identity,
    is_point,
    is_vector,
    kurtosis,
    local_metric,
    local_metric_jacobian,
    log,
    log!,
    log_local_metric_density,
    manifold_dimension,
    metric,
    mean,
    mean!,
    mean_and_var,
    mean_and_std,
    median,
    median!,
    mid_point,
    mid_point!,
    minkowski_metric,
    moment,
    norm,
    normal_tvector_distribution,
    number_eltype,
    one,
    power_dimensions,
    project,
    project!,
    projected_distribution,
    real_dimension,
    ricci_curvature,
    ricci_tensor,
    representation_size,
    retract,
    retract!,
    riemann_tensor,
    set_component!,
    sharp,
    sharp!,
    shortest_geodesic,
    skewness,
    std,
    sym_rem,
    symplectic_inverse_times,
    symplectic_inverse_times!,
    submanifold,
    submanifold_component,
    submanifold_components,
    uniform_distribution,
    var,
    vector_space_dimension,
    vector_transport_along,
    vector_transport_along!,
    vector_transport_direction,
    vector_transport_direction!,
    vector_transport_to,
    vector_transport_to!,
    vee,
    vee!,
    zero_vector,
    zero_vector!
# Lie group types & functions
export AbstractGroupAction,
    AbstractGroupOperation,
    ActionDirection,
    AdditionOperation,
    CircleGroup,
    GeneralLinear,
    GroupManifold,
    GroupOperationAction,
    Identity,
    InvariantMetric,
    LeftAction,
    LeftInvariantMetric,
    MultiplicationOperation,
    ProductGroup,
    ProductOperation,
    RealCircleGroup,
    RightAction,
    RightInvariantMetric,
    RotationAction,
    SemidirectProductGroup,
    SpecialEuclidean,
    SpecialLinear,
    SpecialOrthogonal,
    TranslationGroup,
    TranslationAction
export AbstractInvarianceTrait
export IsMetricManifold, IsConnectionManifold
export IsGroupManifold,
    HasLeftInvariantMetric, HasRightInvariantMetric, HasBiinvariantMetric
export adjoint_action,
    adjoint_action!,
    affine_matrix,
    apply,
    apply!,
    apply_diff,
    apply_diff!,
    base_group,
    center_of_orbit,
    has_approx_invariant_metric,
    compose,
    compose!,
    direction,
    exp_lie,
    exp_lie!,
    group_manifold,
    geodesic,
    get_coordinates_lie,
    get_coordinates_lie!,
    get_coordinates_orthogonal,
    get_coordinates_orthonormal,
    get_coordinates_orthogonal!,
    get_coordinates_orthonormal!,
    get_vector_diagonalizing!,
    get_vector_lie,
    get_vector_lie!,
    get_vector_orthogonal,
    get_vector_orthonormal,
    get_coordinates_vee!,
    has_biinvariant_metric,
    has_invariant_metric,
    identity_element,
    identity_element!,
    inv,
    inv!,
    inverse_apply,
    inverse_apply!,
    inverse_apply_diff,
    inverse_apply_diff!,
    inverse_translate,
    inverse_translate!,
    inverse_translate_diff,
    inverse_translate_diff!,
    lie_bracket,
    lie_bracket!,
    log_lie,
    log_lie!,
    optimal_alignment,
    optimal_alignment!,
    screw_matrix,
    switch_direction,
    translate,
    translate!,
    translate_diff,
    translate_diff!
# Orthonormal bases
export AbstractBasis,
    AbstractOrthonormalBasis,
    DefaultOrthonormalBasis,
    DiagonalizingOrthonormalBasis,
    ProjectedOrthonormalBasis,
    CachedBasis,
    DiagonalizingBasisData,
    ProductBasisData,
    PowerBasisData
export OutOfInjectivityRadiusError
export get_basis,
    get_coordinates, get_coordinates!, get_vector, get_vector!, get_vectors, number_system

# atlases and charts
export get_point, get_point!, get_parameters, get_parameters!

end # module


mutable struct CameraIntrinsic
  K::Array{Float64,2}
  CameraIntrinsic(::Nothing) = new()
  CameraIntrinsic(;x0=320.0,y0=240.0,fx=510.0,fy=510.0,s=0.0) = new([[fx;s;x0]';[0.0;fy;y0]';[0.0;0;1]'])
end

# Camera extrinsic must be world in camera frame (cRw)
mutable struct CameraExtrinsic
  R::SO3
  t::Vector{Float64}
  CameraExtrinsic(::Nothing) = new()
  CameraExtrinsic(;R=SO3(0),t=zeros(3)) = new(R, t)
end
mutable struct CameraModelFull
  ci::CameraIntrinsic
  ce::CameraExtrinsic
  # cd::CameraDistortion
  CameraModelFull(::Nothing) = new()
  CameraModelFull(;ci=CameraIntrinsic(), ce=CameraExtrinsic()) = new(ci,ce)
end
function project!(ret::Vector{Float64}, ci::CameraIntrinsic, ce::CameraExtrinsic, pt::Vector{Float64})
  res = ci.K*(ce.R.R*pt + ce.t)
  ret[1:2] = res[1:2]./res[3]
  nothing
end
project!(ret::Vector{Float64}, cm::CameraModelFull, pt::Vector{Float64}) = project!(ret, cm.ci, cm.ce, pt)
function project(cm::CameraModelFull, pt::Vector{Float64})
  res = Vector{Float64}(2)
  project!(res, cm, pt)
  return res
end

# pinhole camera model
# (x, y)/f = (X, Y)/Z
function cameraResidual!(
      res::Vector{Float64},
      z::Vector{Float64},
      ci::CameraIntrinsic,
      ce::CameraExtrinsic,
      pt::Vector{Float64}  )
  # in place memory operations
  project!(res, ci, ce, pt)
  res[1:2] .*= -1.0
  res[1:2] += z[1:2]
  nothing
end


module AntiIdClosedGroup

using ClosedGroupFunctions
using Zd_Arithmetics
using SparseArrays
using LinearAlgebra
using Combinatorics
using Bijections

import Zd_Arithmetics: ℤ₂, ℤ₄
import Base: inv

export

    id,
    anti_id,

    po2ai,
    ai2po,
    po2matrix,
    po2latex,
    all_pos2latex,
    inv,

    first_po,
    next_po,
    construct_generators,

    perm_rows,
    perm_cols,
    print_permutation_cycles,
    apply_permutation,

    ord,
    tr_of_power,
    rank_of_power_plus_id,
    ℤ₂ⁿ,
    ker_g_plus_id,
    quadratic_form,
    GArf,
    generalized_arf

include("anti-identity.jl")
include("notation_converters.jl")
include("generators.jl")
include("permutations.jl")
include("conjugacy_invariants.jl")

end # module


include("../include.jl")

using JuMP, Ipopt

m = Model()
@variable(m, x[1:3])

eps = 1e-4
scaler = 1.1
@objective(m, Min, sum(x[i] for i = 1:3) )
@NLconstraint(m, scaler * x[1] <= 0)
@NLconstraint(m, x[1] >= eps)
@NLconstraint(m, scaler^2 * x[2] >= 0)
@NLconstraint(m, x[2] <= eps)
@NLconstraint(m, scaler^3 * x[3] >= 0)
@NLconstraint(m, x[3] <= eps)

nlp_raw = MathProgNLPModel(m);
my_pars = Class_parameters();
it = one_phase_solve(nlp_raw,my_pars);

IpoptSolve(nlp_raw)

#it.point.x


#
# Merge Sort
#

function mergeSort(A)

  Anew = copy(A)
  mergesort_array(Anew,  A,  1,  length(A)+1)

end

function mergesort_array(A, result, left, right)

  if right - left < 2
    return
  end
  if right - left == 2
    if result[left] > result[left+1]
      result[left],  result[left+1] = result[left+1], result[left]
    end
    return
  end

  mid = Int(round((right + left) / 2))
  mergesort_array(result, A, left,  mid)
  mergesort_array(result, A, mid, right)

  i = left
  j = mid
  idx = left
  while idx < right
    if j >= right || (i < mid && A[i] < A[j])
      result[idx] = A[i]
      i = i + 1
    else
      result[idx] = A[j]
      j = j + 1
    end
    idx = idx + 1
  end

end

function main()

  # list to be sorted
  rng = MersenneTwister(1234);
  A = shuffle(rng, Vector(1:100))
  Aorigin = copy(A)

  # sorting
  @time mergeSort(A)

  print("The original list: ")
  println(Aorigin)

  print("The sorted list: ")
  println(A)

end

if contains(@__FILE__, PROGRAM_FILE)
    main()
end


#=

 Test cases for interval exceptions from IEEE Std 1788-2015

 Copyright 2016 Oliver Heimlich (oheim@posteo.de)

 Copying and distribution of this file, with or without modification,
 are permitted in any medium without royalty provided the copyright
 notice and this notice are preserved.  This file is offered as-is,
 without any warranty.

=#
#Language imports

#Test library imports
using Test

#Arithmetic library imports
using IntervalArithmetic

#Preamble
setprecision(53)
setprecision(Interval, Float64)
setrounding(Interval, :tight)
# Set full format, and show decorations
@format full

@testset "exceptions" begin
    @test @interval("[+infinity]") == ∅
    @test Interval(Inf, -Inf) == ∅
    @test @interval("[1.0000000000000001, 1.0000000000000002]") == Interval(1.0, 0x1.0000000000001p+0)
end
# FactCheck.exitstatus()


import ClimaCore:
    Fields,
    Domains,
    Topologies,
    Meshes,
    DataLayouts,
    Operators,
    Geometry,
    Spaces

using OrdinaryDiffEq: ODEProblem, solve, SSPRK33

import Logging
import TerminalLoggers
Logging.global_logger(TerminalLoggers.TerminalLogger())

const FT = Float64

n = 32
z₀ = FT(0)
z₁ = FT(10)
t₀ = FT(0)
t₁ = FT(10)
μ = FT(-1 / 2)
ν = FT(5)
𝓌 = FT(1)
δ = FT(1)

domain = Domains.IntervalDomain(
    Geometry.ZPoint{FT}(z₀),
    Geometry.ZPoint{FT}(z₁),
    boundary_names = (:bottom, :top),
)
mesh = Meshes.IntervalMesh(domain, nelems = n)

cs = Spaces.CenterFiniteDifferenceSpace(mesh)
fs = Spaces.FaceFiniteDifferenceSpace(cs)
zc = Fields.coordinate_field(cs)
zp = (z₀ + z₁ / n / 2):(z₁ / n):(z₁ - z₁ / n / 2)

function gaussian(z, t; μ = -1 // 2, ν = 1, 𝓌 = 1, δ = 1)
    return exp(-(z - μ - 𝓌 * t)^2 / (4 * ν * (t + δ))) / sqrt(1 + t / δ)
end
function ∇gaussian(z, t; μ = -1 // 2, ν = 1, 𝓌 = 1, δ = 1)
    return -2 * (z - μ - 𝓌 * t) / (4 * ν * (δ + t)) *
           exp(-(z - μ - 𝓌 * t)^2 / (4 * ν * (δ + t))) / sqrt(1 + t / δ)
end

T = gaussian.(zc.z, -0; μ = μ, δ = δ, ν = ν, 𝓌 = 𝓌)
V = Geometry.WVector.(ones(FT, fs))

# Solve Adv-Diff Equation: ∂_t T = α ∇²T
z₀ = FT(0)
z₁ = FT(10)

function ∑tendencies!(dT, T, z, t)

    ic2f = Operators.InterpolateC2F()
    bc_vb = Operators.SetValue(FT(gaussian(z₀, t; ν = ν, δ = δ, 𝓌 = 𝓌, μ = μ)))
    bc_vt = Operators.SetValue(FT(gaussian(z₁, t; ν = ν, δ = δ, 𝓌 = 𝓌, μ = μ)))
    bc_gb = Operators.SetGradient(
        Geometry.WVector(FT(∇gaussian(z₀, t; ν = ν, δ = δ, 𝓌 = 𝓌, μ = μ))),
    )
    bc_gt = Operators.SetGradient(
        Geometry.WVector(FT(∇gaussian(z₁, t; ν = ν, δ = δ, 𝓌 = 𝓌, μ = μ))),
    )

    #   Upwind Biased Product
    #   UB = Operators.UpwindBiasedProductC2F(
    #       bottom = Operators.Extrapolate(),
    #       top = bc_vt,
    #   )
    #   ∂ = Operators.GradientF2C()
    #   return @. dT = -∂(UB(V, ic2f(T)))

    A = Operators.AdvectionC2C(bottom = bc_vb, top = Operators.Extrapolate())

    gradc2f = Operators.GradientC2F(bottom = bc_vb, top = bc_gt)
    divf2c = Operators.DivergenceF2C()

    return @. dT = divf2c(ν * gradc2f(T)) - A(V, T)
end

@show ∑tendencies!(similar(T), T, nothing, 0.0);

# Solve the ODE operator
Δt = 0.0001

prob = ODEProblem(∑tendencies!, T, (t₀, t₁))
sol = solve(
    prob,
    SSPRK33(),
    dt = Δt,
    saveat = 10000 * Δt,
    progress = true,
    progress_message = (dt, u, p, t) -> t,
);

ENV["GKSwstype"] = "nul"
using ClimaCorePlots, Plots
Plots.GRBackend()

dirname = "advect_diffusion"
path = joinpath(@__DIR__, "output", dirname)
mkpath(path)

anim = Plots.@animate for (nt, u) in enumerate(sol.u)
    Plots.plot(
        u,
        xlim = (0, 1),
        ylim = (-1, 10),
        title = "$(nt-1) s",
        lc = :black,
        lw = 2,
        ls = :dash,
        label = "Approx Sol.",
        legend = :outerright,
        m = :o,
        xlabel = "T(z)",
        ylabel = "z",
    )
    Plots.plot!(
        gaussian.(zp, nt - 1; μ = μ, δ = δ, ν = ν, 𝓌 = 𝓌),
        zp,
        xlim = (0, 1),
        ylim = (-1, 10),
        title = "$(nt) s",
        lc = :red,
        lw = 2,
        label = "Analytical Sol.",
        m = :x,
    )
end
Plots.mp4(anim, joinpath(path, "advect_diffusion.mp4"), fps = 10)
Plots.png(
    Plots.plot(sol.u[end], xlim = (0, 1)),
    joinpath(path, "advect_diffusion_end.png"),
)

function linkfig(figpath, alt = "")
    # buildkite-agent upload figpath
    # link figure in logs if we are running on CI
    if get(ENV, "BUILDKITE", "") == "true"
        artifact_url = "artifact://$figpath"
        print("\033]1338;url='$(artifact_url)';alt='$(alt)'\a\n")
    end
end

linkfig(
    relpath(
        joinpath(path, "advect_diffusion_end.png"),
        joinpath(@__DIR__, "../.."),
    ),
    "Advection-Diffusion End Simulation",
)


#
# This file is a part of MolecularGraph.jl
# Licensed under the MIT License http://opensource.org/licenses/MIT
#

export
    combinations, sortstablemax, sortstablemin

## combinations

struct Combinations
    n::Int
    k::Int
    indice::Vector{Int}
end

function Base.iterate(cmb::Combinations, state=nothing)
    if cmb.indice[end] == cmb.n
        cmb.k == 1 && return
        i = 1
        while cmb.indice[end - i] == cmb.n - i
            i += 1
            i == cmb.k && return
        end
        cmb.indice[cmb.k - i] += 1
        for j in (cmb.k - i + 1):cmb.k
            cmb.indice[j] = cmb.indice[j - 1] + 1
        end
    else
        cmb.indice[end] += 1
    end
    return (cmb.indice, state)
end

Base.eltype(::Type{<:Combinations}) = Vector{Int}
Base.IteratorSize(::Type{<:Combinations}) = Base.SizeUnknown()

"""
    combinations(n::Int, k::Int=2)

An iterator that yields ``k``-combinations (cardinality ``k`` subsets of the cardinality ``n`` set) as an array of integers.
"""
function combinations(n::Int, k::Int=2)
    if k < 0
        throw(DomainError(k, "k should not be negative"))
    elseif k > n
        throw(ErrorException("$(k) should be equal or less than the collection size"))
    elseif k == 0
        return (Int[],)
    else
        idc = collect(1:k)
        idc[end] -= 1
        return Combinations(n, k, idc)
    end
end

function sortstablemax(iter; by=identity, kwargs...)
    isempty(iter) && return kwargs[:init]
    cmp(x, y) = by(x) < by(y) ? y : x
    return reduce(cmp, iter)
end

function sortstablemin(iter; by=identity, kwargs...)
    isempty(iter) && return kwargs[:init]
    cmp(x, y) = by(x) > by(y) ? y : x
    return reduce(cmp, iter)
end


@testset "81.search-in-rotated-sorted-array-ii.jl" begin
    @test search_rotate([2, 5, 6, 0, 0, 1, 2], 0) == true
    @test search_rotate([2, 5, 6, 0, 0, 1, 2], 3) == false
    @test search_rotate([2, 5, 6, 0, 0, 1, 2], 2) == true
    @test search_rotate([2, 5, 6, 0, 0, 1, 2], 1) == true
    @test search_rotate([2, 5, 6, 0, 0, 1, 2], 4) == false
    @test search_rotate([2, 5, 6, 0, 0, 1, 2], 5) == true
    @test search_rotate([2, 5, 6, 0, 0, 1, 2], 6) == true
    @test search_rotate([2, 5, 6, 0, 0, 1, 2], 7) == false
    arr = collect(1:30)
    for i in 1:29
        na = circshift(arr, i)
        for j in 1:30
            @test search_rotate(na, j) == true
        end
    end
end


#MDDatasets: Additional array/vector operations
#-------------------------------------------------------------------------------
#=NOTE:
These tools should eventually be moved to a separate unit.
=#

#==Iterators
===============================================================================#
#Obtain an array iterator:
subscripts(sz::Tuple) = [Tuple(x) for x in CartesianIndices(sz)]
subscripts(a::Array) = [Tuple(x) for x in CartesianIndices(a)]

#==Useful tests
===============================================================================#

#Verifies that v is strictly increasing (no repeating values):
function isincreasing(v::Vector)
	prev = v[1]
	for x in v[2:end]
		if x <= prev
			return false
		end
	end
	return true
end

isincreasing(r::AbstractRange) = (step(r) > 0)

#==Validation functions (throw exceptions)
===============================================================================#

#==Helper functions
===============================================================================#

#Assumes vector is ordered... Not currently checking that it is true..
function findclosestindex(v::Vector, val)
	reltol =  1/1000 #WANTCONST
	if length(v) < 2
		if abs(1-val/v[1]) < reltol
			return 1
		else
			throw("Value not found: $val")
		end
	end
	vlast = v[2] #Gets an order of magnitude for first point
	for idx in 1:length(v)
		Δ = abs(v[idx] - vlast)
		if abs(val-v[idx]) < reltol*Δ
			return idx
		end
		vlast = v[idx]
	end
	throw("Value not found: $val")
end

#Create vector with data left-shifted by n (padded with 0s)
function _lshift(v::Vector, n::Int)
	@assert(n >= 0, "Cannot shfit by a negative number")
	result = zeros(v)
	for i in 1:(length(v)-n)
		result[i] = v[i+n]
	end
	return result
end

#Create vector with data right-shifted by n (padded with 0s)
function _rshift(v::Vector, n::Int)
	@assert(n >= 0, "Cannot shfit by a negative number")
	result = zero(v)
	for i in length(v):-1:(1+n)
		result[i] = v[i-n]
	end
	return result
end

#Create vector with data shifted by +/-n (padded with 0s)
function shift(v::Vector, n::Int)
	if n >= 0
		return _rshift(v, n)
	else
		return _lshift(v, -n)
	end
end

#Compute difference between two adjacent points:
#TODO: optimize operations so they run faster
function delta(v::Vector)
	return v[2:end] .- v[1:end-1]
end

#Compute mean of two adjacent points:
#TODO: optimize operations so they run faster
function meanadj(v::Vector)
	return (v[1:end-1] .+ v[2:end])./2
end

#Last line


using Revise
using ADCME
using PyPlot
matplotlib.use("macosx")
X,_,_,_ = mnist()
X = X/127.5 .-1.
Y = zeros(size(X,1), 28*28)
for i = 1:size(X,1)
    Y[i,:] = vec(X[i,:,:,1])
end
X = copy(Y)

function discriminator(x, gan)
    y = ae(x, [256,128,64,1])
    y = sigmoid(y)
    return y
end

function generator(x, gan)
    y = ae(x, [128,256,784])
    return y
end

# function loss(t, f, gan)
#     t_logits = discriminator(t, gan)
#     f_logits = discriminator(f, gan)
#     t_loss = mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=t_logits, labels=tf.ones_like(t_logits)))
#     f_loss = mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=f_logits, labels=tf.zeros_like(f_logits)))
#     d_loss = t_loss + f_loss
#     g_loss = mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=f_logits, labels=tf.ones_like(f_logits)))
#     d_loss, g_loss
# end

reset_default_graph()
gan = GAN(X, generator, discriminator, "klgan"; latent_dim=100, batch_size=32)
opt_d = RMSPropOptimizer(5e-5).minimize(gan.d_loss, var_list=gan.d_vars)
opt_g = RMSPropOptimizer(5e-5).minimize(gan.g_loss, var_list=gan.g_vars)
img = sample(gan, 25)

# if gantype=="wgan"
#     op = clip(gan.d_vars[1:2:end], -0.001, 0.001)
#     [push!(gan.update, o) for o in op]
# end
sess = Session(); init(sess)
rm("images/", force=true, recursive=true)
mkdir("images/")
for i = 1:200000
    feed = Dict(gan.noise=>randn(gan.batch_size, gan.latent_dim), gan.ids=>Int32.(rand(1:size(X,1), gan.batch_size)))
    _, dl, _ = run(sess, [opt_d, gan.d_loss,gan.update], feed_dict=feed)
    for i = 1:4
        run(sess, opt_g, feed_dict=feed)
    end
    _, gl = run(sess, [opt_g, gan.g_loss], feed_dict=feed)
    if i%100==1
        println("iter=$i, dloss=$dl, gloss=$gl")
    end
    if i%500==1
        val = run(sess, img)
        fig, axs = subplots(5,5)
        cnt = 1
        for i = 1:5
            for j = 1:5
                axs[i,j].imshow(reshape(val[cnt,:], 28, 28), cmap="gray")
                axs[i,j].axis("off")
                cnt += 1
            end
        end
        savefig("images/$i.png")
        close("all")
    end
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENCE in the project root.
# ------------------------------------------------------------------

"""
    LWR(var₁=>param₁, var₂=>param₂, ...)

Locally weighted regression estimation solver.

## Parameters

* `neighbors` - Number of neighbors (default to all the data)
* `distance`  - A distance from Distances.jl (default to `Euclidean()`)
* `weightfun` - Weighting function (default to `exp(-3*h^2/2)`)

### References

* Stone 1977. *Consistent non-parametric regression.*
* Cleveland 1979. *Robust locally weighted regression and smoothing scatterplots.*
* Cleveland & Grosse 1991. *Computational methods for local regression.*
"""
@estimsolver LWR begin
  @param neighbors = nothing
  @param distance = Euclidean()
  @param weightfun = h -> exp(-3*h^2)
end

function solve(problem::EstimationProblem, solver::LWR)
  # retrieve problem info
  pdata = data(problem)
  pdomain = domain(problem)

  mactypeof = Dict(name(v) => mactype(v) for v in variables(problem))

  # result for each variable
  μs = []; σs = []

  for covars in covariables(problem, solver)
    for var in covars.names
      # get user parameters
      varparams = covars.params[(var,)]

      # determine value type
      V = mactypeof[var]

      # retrieve non-missing data
      locs = findall(!ismissing, pdata[var])
      𝒟 = view(pdata, locs)
      n = nelements(𝒟)

      # determine number of nearest neighbors to use
      k = isnothing(varparams.neighbors) ? n : varparams.neighbors

      # determine distance type
      D = varparams.distance

      # weight function
      w = varparams.weightfun

      @assert n > 0 "estimation requires data"

      @assert k ≤ n "invalid number of neighbors"

      # fit search tree
      X = [coordinates(centroid(𝒟, i)) for i in 1:nelements(𝒟)]
      if D isa NearestNeighbors.MinkowskiMetric
        tree = KDTree(X, D)
      else
        tree = BallTree(X, D)
      end

      # lookup non-missing values
      z = 𝒟[var]

      # pre-allocate memory for results
      varμ = Vector{V}(undef, nelements(pdomain))
      varσ = Vector{V}(undef, nelements(pdomain))

      # estimation loop
      for loc in traverse(pdomain, LinearPath())
        x = coordinates(centroid(pdomain, loc))

        # find neighbors
        is, ds = knn(tree, x, k)
        δs = ds ./ maximum(ds)

        # weighted least-squares
        Wₗ = Diagonal(w.(δs))
        Xₗ = [ones(eltype(x), k) reduce(hcat, X[is])']
        zₗ = view(z, is)
        θₗ = Xₗ'*Wₗ*Xₗ \ Xₗ'*Wₗ*zₗ

        # linear combination of response values
        xₒ = [one(eltype(x)); x]
        ẑₒ = θₗ ⋅ xₒ
        rₗ = Wₗ*Xₗ*(Xₗ'*Wₗ*Xₗ\xₒ)
        r̂ₒ = norm(rₗ)

        varμ[loc] = ẑₒ
        varσ[loc] = r̂ₒ
      end

      push!(μs, var => varμ)
      push!(σs, Symbol(var,"_variance") => varσ)
    end
  end

  georef((; μs..., σs...), pdomain)
end


using Pkg
Pkg.activate(".") # Create new environment in this folder

# first time you need to install dependencies
Pkg.add("Distributions")
Pkg.add("StatsBase")
Pkg.add(["DataFrames","DataFramesMeta","Chain"])
Pkg.add("Plots")
Pkg.add("CategoricalArrays")
Pkg.add("LightGraphs")
Pkg.add("FixedEffectModels")

# past the first time, you only need to instanciate the current folder
Pkg.instantiate() # Updates packages given .toml file

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

using Distributions
using LinearAlgebra
using StatsBase
using DataFrames
using Plots
using CategoricalArrays
using FixedEffectModels
using Chain
using DataFramesMeta
using LightGraphs
using TikzGraphs

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

α_sd = 1
ψ_sd = 1

csort = 0.5 # Sorting effect
csig  = 0.5 # Cross-sectional standard deviation
w_sigma = 0.2
λ = 0.1 # Moving Probability is fixed
nt = 10

function data_generating_process(α_sd,
                                ψ_sd,
                                csort,
                                csig,
                                w_sigma,
                                λ,
                                nt)

    cnetw = 0.2 # Network effect

    nk = 30
    nl = 10

    # approximate each distribution with some points of support
    ψ = quantile.(Normal(), (1:nk) / (nk + 1)) * α_sd
    α = quantile.(Normal(), (1:nl) / (nl + 1)) * ψ_sd

    # Let's create type-specific transition matrices
    # We are going to use joint normals centered on different values
    G = zeros(nl, nk, nk)
    for l in 1:nl, k in 1:nk
        G[l, k, :] = pdf( Normal(0, csig), ψ .- cnetw * ψ[k] .- csort * α[l])
        G[l, k, :] = G[l, k, :] ./ sum(G[l, k, :])
    end

    # We then solve for the stationary distribution over psis for each alpha value
    # We apply a crude fixed point approach
    H = ones(nl, nk) ./ nk
    for l in 1:nl
        M = transpose(G[l, :, :])
        for i in 1:100
            H[l, :] = M * H[l, :]
        end
    end

    # p1 = plot(G[1, :, :], xlabel="Previous Firm", ylabel="Next Firm", zlabel="G[1, :, :]", st=:wireframe)
    # p2 = plot(G[nl, :, :], xlabel="Previous Firm", ylabel="Next Firm", zlabel="G[nl, :, :]", st=:wireframe, right_margin = 10Plots.mm) # right_margin makes sure the figure isn't cut off on the right
    # plot(p1, p2, layout = (1, 2), size=[600,300])

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

    nt = 10
    ni = 10000

    # We simulate a balanced panel
    ll = zeros(Int64, ni, nt) # Worker type
    kk = zeros(Int64, ni, nt) # Firm type
    spellcount = zeros(Int64, ni, nt) # Employment spell

    for i in 1:ni

        # We draw the worker type
        l = rand(1:nl)
        ll[i,:] .= l

        # At time 1, we draw from H
        kk[i,1] = sample(1:nk, Weights(H[l, :]))

        for t in 2:nt
            if rand() < λ
                kk[i,t] = sample(1:nk, Weights(G[l, kk[i,t-1], :]))
                spellcount[i,t] = spellcount[i,t-1] + 1
            else
                kk[i,t] = kk[i,t-1]
                spellcount[i,t] = spellcount[i,t-1]
            end
        end

    end

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

    firms_per_type = 15
    jj = zeros(Int64, ni, nt) # Firm identifiers

    draw_firm_from_type(k) = sample(1:firms_per_type) + (k - 1) * firms_per_type

    for i in 1:ni

        # extract firm type
        k = kk[i,1]

        # We draw the firm (one of firms_per_type in given group)
        jj[i,1] = draw_firm_from_type(k)

        for t in 2:nt
            if spellcount[i,t] == spellcount[i,t-1]
                # We keep the firm the same
                jj[i,t] = jj[i,t-1]
            else
                # We draw a new firm
                k = kk[i,t]

                new_j = draw_firm_from_type(k)
                # Make sure the new firm is actually new
                while new_j == jj[i,t-1]
                    new_j = draw_firm_from_type(k)
                end

                jj[i,t] = new_j
            end
        end
    end
    # Make sure firm ids are contiguous
    contiguous_ids = Dict( unique(jj) .=> 1:length(unique(jj))  )
    jj .= getindex.(Ref(contiguous_ids),jj);

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

    ii = repeat(1:ni,1,nt)
    tt = repeat((1:nt)',ni,1)
    df = DataFrame(i=ii[:], j=jj[:], l=ll[:], k=kk[:], α=α[ll[:]], ψ=ψ[kk[:]], t=tt[:], spell=spellcount[:]);

    return df, G, H

end

#Number of workers in the firm:

df,G,H = data_generating_process(α_sd,ψ_sd,csort,csig, w_sigma, λ, nt);

crossSectionMean = @chain df begin
    groupby([:j, :t])
    combine(nrow => :count)
    groupby(:j)
    combine(:count => mean)
    @aside println("On average each firm has $(round(mean(_.count_mean))) observations, averaging across time periods.")
    # first(_, 5)
end

#Number of movers at the firm, defined as people that moved into the firm:

println("Maximum number of moves by a single worker: ", maximum(df.spell))

moversDataFrame = @chain df begin
    groupby([:j, :i, :spell])
    combine(:t => mean)
    @transform!(:spell_true = :spell .> 0)
    groupby(:j)
    combine(:spell_true => sum)
    @aside println("On average each firm has $(round(mean(_.spell_true_sum))) movers over the 10 periods.")
end

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

df[!, :lw] = df.α + df.ψ + w_sigma * rand(Normal(), size(df)[1]);
df[!, :dummy] .= 1;

#1. Mean wage within firm:
eventStudyPanel =  @chain df begin
    groupby([:j, :t]) # Group by firm and time
    transform(:lw => mean) # Obtain wage at that group level
    @transform!(:wage_percentile = cut(:lw_mean,4)) # Get wage_percentile at same level of agregation
    groupby([:i, :spell])
    transform(nrow => :years_at_firm) # Get number of years individual spend at a single firm
    sort([:i,:t], rev = true) # Sort over time by individual (this helps to get first 2 periods of last firm)
    groupby([:i, :spell])  # by group and spell
    transform(:dummy .=>  cumsum => :years_at_firm_sofar) # Get the number of years spent at a firm so far
    sort([:i,:t], rev = false) # Sort over time by individual (this helps to get first 2 periods of last firm)
    groupby([:i, :spell])  # by group and spell
    transform(:dummy .=>  cumsum => :years_at_firm_sofar_inverse) # Get the number of years spent at a firm so far, but backwards.
    @aside initialFirmDataFrame = @chain _ begin
        subset(:spell => ByRow(==(1)), :years_at_firm_sofar_inverse => ByRow(<=(2))) # Generate dataframe for initial firm, keep last two years (call it initialFirmDataFrame)
    end
    subset(:spell => ByRow(==(0)), :years_at_firm_sofar => ByRow(<=(2))) # Generate dataframe for subsequent firm, keep first two years
    append!(initialFirmDataFrame) # Append both dataframes
    groupby(:i) # group by person
    transform(nrow => :nreps) # and get number of repetitions
    subset(:nreps => ByRow(==(4))) # and get number of repetitions
    sort([:i,:t]) # sort by i and t to check if it worked out
    # Generating event time variable:
    sort([:i,:t], rev = false) # Sort over time by individual (this helps to get first 2 periods of last firm)
    groupby([:i])  # by worker and time
    transform(:dummy .=>  cumsum => :event_time) # Get the number of years spent at a firm so far, but backwards.
end

# Generate dictionary to index cut label:
percentile_cut = Dict(1:4 .=> unique(sort(eventStudyPanel.wage_percentile)))

# Get Card Heining and Kling event Figure:
# There may be a smart way to do this, here I'm brute forcing it ...
initial = 1
final = 4
function generateEventStudy(eventStudyPanel, initial, final)
    return  eventStudy =  @chain eventStudyPanel begin
                            @aside finalJob = @chain _ begin
                                subset(:wage_percentile => ByRow(==(percentile_cut[final])), :spell => ByRow(==(1)))
                            end
                            subset(:wage_percentile => ByRow(==(percentile_cut[initial])), :spell => ByRow(==(0)))
                            append!(finalJob) # Append both dataframes
                            groupby(:i)
                            transform(nrow => :nreps)
                            subset(:nreps => ByRow(==(4)))
                            groupby([:wage_percentile, :event_time])
                            combine(:lw => mean)
                            sort(:event_time)
    end
end

eventStudySwitchers = [generateEventStudy(eventStudyPanel, 1, 1).lw_mean,
                        generateEventStudy(eventStudyPanel, 1, 2).lw_mean,
                        generateEventStudy(eventStudyPanel, 1, 3).lw_mean,
                        generateEventStudy(eventStudyPanel, 1, 4).lw_mean]
eventStudySwitchers2 = [generateEventStudy(eventStudyPanel, 4, 1).lw_mean,
                        generateEventStudy(eventStudyPanel, 4, 2).lw_mean,
                        generateEventStudy(eventStudyPanel, 4, 3).lw_mean,
                        generateEventStudy(eventStudyPanel, 4, 4).lw_mean];

plot(1:4,eventStudySwitchers, label=["1 to 1" "1 to 2" "1 to 3" "1 to 4"],markershape = :square)
plot!(1:4,eventStudySwitchers2,label =  ["4 to 1 " "4 to 2" "4 to 3" "4 to 4"],markershape = :circle)
plot!(legend=:outertopright)

# Calibrating the parameters:
#----------------------------------------------------------------------------------------

# Generate function with variance decomposition:

function variance_decomposition(df)
    varianceDecomposition =  @chain df begin
        groupby([:k]) # Group by firm type
        combine(:α => mean, :ψ => mean)
    end

    std_α = std(varianceDecomposition.α_mean)
    std_ψ = std(varianceDecomposition.ψ_mean)
    std_αψ = cov(varianceDecomposition.α_mean,varianceDecomposition.ψ_mean)

    return [std_α std_ψ std_αψ]
end

α_sd = 1
ψ_sd = 1
csort = 0.5 # Sorting effect
csig  = 0.5 # Cross-sectional standard deviation
w_sigma = 0.2

gap = 0.20
α_sd_list = 0.01:gap:1
ψ_sd_list = 0.01:gap:1
csort_list = 0.01:gap:1 # Sorting effect
csig_list  = 0.01:gap:1 # Cross-sectional standard deviation
w_sigma_list = 0.01:gap:1

flat_gridpoints(grids) = vec(collect(Iterators.product(grids...)))
grid_points = flat_gridpoints((α_sd_list, ψ_sd_list, csort_list, csig_list,w_sigma_list));

jj = 1

true_var_decomp = [0.186 0.101 0.110]
difference_outcome = []

#Apply grid search, this can be optimized for sure...
for jj in 1:length(grid_points)
    α_sd,ψ_sd,csort,csig, w_sigma = grid_points[jj]
    df,_,_ = data_generating_process(α_sd,ψ_sd,csort,csig, w_sigma, λ, nt);
    var_decomp = variance_decomposition(df)
    euclidean_dist = sum((true_var_decomp-var_decomp).^2)
    push!(difference_outcome, euclidean_dist)
end

# Get index of minimum euclidean distance estimation:
_, index_calibration = findmin(difference_outcome)
α_sd,ψ_sd,csort,csig, w_sigma = grid_points[index_calibration]
df, G, H  = data_generating_process(α_sd,ψ_sd,csort,csig, w_sigma, λ, nt);

p1 = plot(G[1, :, :], xlabel="Previous Firm", ylabel="Next Firm", zlabel="G[1, :, :]", st=:wireframe)
p2 = plot(G[nl, :, :], xlabel="Previous Firm", ylabel="Next Firm", zlabel="G[nl, :, :]", st=:wireframe, right_margin = 10Plots.mm) # right_margin makes sure the figure isn't cut off on the right
plot(p1, p2, layout = (1, 2), size=[600,300])

plot(H, xlabel="Worker", ylabel="Firm", zlabel="H", st=:wireframe)

# The output reduces cross sectional variance too much, check what happens when ↑ csig ...
df, G, H  = data_generating_process(α_sd,ψ_sd,csort, .1, w_sigma, λ, nt);
var_decomp = variance_decomposition(df)
euclidean_dist = sum((true_var_decomp-var_decomp).^2)
println("Euclidean distance is: $euclidean_dist")

p1 = plot(G[1, :, :], xlabel="Previous Firm", ylabel="Next Firm", zlabel="G[1, :, :]", st=:wireframe)
p2 = plot(G[nl, :, :], xlabel="Previous Firm", ylabel="Next Firm", zlabel="G[nl, :, :]", st=:wireframe, right_margin = 10Plots.mm) # right_margin makes sure the figure isn't cut off on the right
plot(p1, p2, layout = (1, 2), size=[600,300])

# It seems csig contributes importantly to covariance! Check this again and get feedback from prof.
# I'll leave it at 0.05 which is 5 times it optimal* value.

# Delete method in case of mistake
# m = @which data_generating_process(α_sd,ψ_sd,csort,csig, w_sigma, λ, nt)
# Base.delete_method(m)

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

# Getting connected set:

function individualDeterministicTransitionMatrix(df,ii)

    nfirms = length(unique(df.j))
    shift_i = @chain df begin
                        subset(:i => ByRow(.==(ii)))
                        sort(:t)
                        combine(first, groupby(_,:spell))
                    end

    nshifts = size(shift_i)[1]
    transitionMatrix_i = zeros(Int32, nfirms, nfirms);

    for ii in 1:nshifts
        if ii != nshifts
            current = shift_i.j[ii]
            next = shift_i.j[ii+1]
            transitionMatrix_i[current,next] = 1
        end
    end

    return transitionMatrix_i
end

function getConnectedDataSet(df)

    # Add all shifts happening in the economy during the whole time
    nfirms = length(unique(df.j))

    totalDeterministicShifts = zeros(nfirms, nfirms);

    for ii in unique(df.i)
        totalDeterministicShifts = totalDeterministicShifts + individualDeterministicTransitionMatrix(df,ii)
    end

    adjacencyMatrix = (UpperTriangular(totalDeterministicShifts) + transpose(LowerTriangular(totalDeterministicShifts))).>1

    adjacencyMatrix = adjacencyMatrix + transpose(adjacencyMatrix)

    println("There are $(sum(totalDeterministicShifts)) job shifts across time ...")

    println("There are $(sum(adjacencyMatrix)) edges between nodes ...")

    simpleGraph = SimpleGraph(adjacencyMatrix);

    connectedNetwork = connected_components(simpleGraph)

    connectedSet = connectedNetwork[1]

    println("We have only $(length(connectedSet)) firms fully connected $(length(connectedSet)/nfirms) of the market...") # Double check we might be wrong...

    return df_connected = df[in(connectedSet).(df.j),:]

    println("Due to unconnectedness we eliminated $(size(df)[1]-size(df_connected)[1]) observations, not much")

end

# Estimate AKM model:
"""
This part of the problem set is for you to implement the AKM estimator.
    As discussed in class, this can be done simply by updating, in turn,
    the worker FE and the firm FE.

    Start by appending 2 new columns `alpha_hat` and `psi_hat` to your data.
    Then loop over the following:

    1. Update `alpha_hat` by taking the mean within `i` net of firm FE
    2. Update `psi_hat` by taking the mean within `fid` net of worker FE

    Question 7:

     - Run the previous steps in a loop, and at each step evaluate how much
     the total mean square error has changed.
     Check that is goes down with every step.
     Stop when the MSE decreases by less than 1e-9.

    Note that you can increase speed by focusing on movers only first.
"""
#--------------------------------------------------------------------------------------

function akm_estimation(df_connected)

    df_connected[:,:alpha_hat] .= .0;
    df_connected[:,:psi_hat] .= .0;

    # Compute firm type fixed effects ols model:
    delta = Inf
    tol = 0.00001
    nIter = 1
    msePast = 0
    while delta>tol

        if nIter == 1
            # Regress controling for industry fixed effects...
            model_j = reg(df_connected, term(:lw) ~ fe(:j), save=true);
            # Obtain residuals (which are controled by industry fe):
            df_connected[:,:alpha_hat] = residuals(model_j);
        else
            # Just obtain the alpha parameters (not averaged) by netting out psi_hat from prev iter.
            df_connected[:,:alpha_hat] = df_connected[:,:lw] - df_connected[:,:psi_hat]
        end

        # Average fixed effects by individual:
        df_connected = transform(groupby(df_connected, :i), :alpha_hat => mean => :alpha_hat)

        # Net out individual fixed effects
        df_connected[:,:psi_hat] = df_connected[:,:lw] - df_connected[:,:alpha_hat]

        # Compute average industry fixed effects
        df_connected = transform(groupby(df_connected, :j), :psi_hat => mean => :psi_hat)

        # Model verification
        model_verify = reg(df_connected, @formula(lw ~ alpha_hat + psi_hat), save=true);

        # Compute residuals and mse
        delta = abs(msePast - sum(residuals(model_verify, df_connected).^2))

        msePast = sum(residuals(model_verify, df_connected).^2)

        # Add count
        nIter = nIter + 1

        println("At iteration number $(nIter), the MSE is $(delta)")

    end

    return df_connected

end

df_connected = getConnectedDataSet(df)

df_connected_results = akm_estimation(df_connected)

# %%

# Limited mobility bias:
#---------------------------------------------------------------------

λ_list = [0.1 0.2 0.3 0.4 0.5 0.6]

nt_list = [5 6 8 10 15]

df,_,_ = data_generating_process(α_sd, ψ_sd, csort, csig, w_sigma, λ, nt);

df_connected = getConnectedDataSet(df)

df_connected_results = akm_estimation(df_connected)


#=
real symmetric or complex Hermitian positive semidefinite matrices (svec scaled triangle)
W ⪰ 0
self-dual
=#

mutable struct PosSemidefTri{C <: RealCompF} <: Cache
    oa_s::Vector{AE}
    d::Int
    w_temp::Vector{RealF}
    W_temp::Matrix{C}
    PosSemidefTri{C}() where {C <: RealCompF} = new{C}()
end

function MOIPajarito.Cones.create_cache(
    oa_s::Vector{AE},
    cone::Hypatia.PosSemidefTriCone{RealF, C},
    ::Optimizer,
) where {C <: RealCompF}
    cache = PosSemidefTri{C}()
    cache.oa_s = oa_s
    dim = MOI.dimension(cone)
    d = cache.d = svec_side(C, dim)
    cache.w_temp = zeros(RealF, dim)
    cache.W_temp = zeros(C, d, d)
    return cache
end

function MOIPajarito.Cones.add_init_cuts(cache::PosSemidefTri{C}, opt::Optimizer) where {C}
    # add variable bounds W_ii ≥ 0
    d = cache.d
    w = cache.oa_s
    JuMP.@constraint(opt.oa_model, [i in 1:d], w[svec_idx(C, i, i)] >= 0)
    opt.use_init_fixed_oa || return

    # add cuts (1, 1, ±2) on (W_ii, W_jj, W_ij), ∀i != j
    # (a linearization of W_ii * W_jj ≥ W_ij^2)
    # initial OA polyhedron is the dual cone of diagonally dominant matrices
    for j in 1:d, i in 1:(j - 1)
        _add_init_cuts(i, j, cache, opt)
    end
    return
end

# real: cuts on (w_ii, w_jj, w_ij) are (1, 1, ±rt2), ∀i != j
function _add_init_cuts(i::Int, j::Int, cache::PosSemidefTri{RealF}, opt::Optimizer)
    w = cache.oa_s
    w_ii = w[svec_idx(RealF, i, i)]
    w_jj = w[svec_idx(RealF, j, j)]
    w_ij = w[svec_idx(RealF, i, j)]
    JuMP.@constraints(opt.oa_model, begin
        w_ii + w_jj + rt2 * w_ij >= 0
        w_ii + w_jj - rt2 * w_ij >= 0
    end)
    return
end

# complex: cuts on smat (w_ii, w_jj, wr_ij, wi_ij) are (1, 1, ±1, ±1), ∀i != j
function _add_init_cuts(i::Int, j::Int, cache::PosSemidefTri{CompF}, opt::Optimizer)
    w = cache.oa_s
    w_ii = w[svec_idx(CompF, i, i)]
    w_jj = w[svec_idx(CompF, j, j)]
    ij_idx = svec_idx(CompF, i, j)
    wr_ij = w[ij_idx]
    wi_ij = w[ij_idx + 1]
    JuMP.@constraints(opt.oa_model, begin
        w_ii + w_jj + wr_ij + wi_ij >= 0
        w_ii + w_jj + wr_ij - wi_ij >= 0
        w_ii + w_jj - wr_ij + wi_ij >= 0
        w_ii + w_jj - wr_ij - wi_ij >= 0
    end)
    return
end

function MOIPajarito.Cones.get_subp_cuts(
    z::Vector{RealF},
    cache::PosSemidefTri,
    opt::Optimizer,
)
    # strengthened cuts from eigendecomposition are λᵢ * rᵢ * rᵢ'
    R = cache.W_temp
    svec_to_smat!(R, z, rt2)
    F = eigen(Hermitian(R, :U), 1e-8, Inf) # TODO tune
    isempty(F.values) && return AE[]
    R_eig = F.vectors * Diagonal(sqrt.(F.values))
    return _get_psd_cuts(R_eig, cache.oa_s, cache, opt)
end

function MOIPajarito.Cones.get_sep_cuts(
    s::Vector{RealF},
    cache::PosSemidefTri,
    opt::Optimizer,
)
    Ws = cache.W_temp
    svec_to_smat!(Ws, s, rt2)
    F = eigen(Hermitian(Ws, :U), -Inf, -opt.tol_feas)
    isempty(F.values) && return AE[]
    return _get_psd_cuts(F.vectors, cache.oa_s, cache, opt)
end


function rescale_velocity!(mdsim::MDSimulation)
	P = mdsim.params
	λ = sqrt(P.d * P.k_B * (P.N - 1) * P.T / (P.m * sum(mdsim.state.v.^2)))
	mdsim.state.v *= λ
end

function update_history!(mdsim::MDSimulation, n::Int)
	mdsim.history.x[:, :, n] = mdsim.state.x
	mdsim.history.v[:, :, n] = mdsim.state.v
	mdsim.history.E_kin[n]   = mdsim.state.E_kin
	mdsim.history.E_pot[n]   = mdsim.state.E_pot
end

function velocity_verlet!(mdsim::MDSimulation)
	P = mdsim.params
	for n in 1:P.n_iter
		@. mdsim.state.x += P.τ * mdsim.state.v + (P.τ^2 / 2) * mdsim.state.a
		compute_accelerations!(mdsim)
		@. mdsim.state.v += (P.τ / 2) * (mdsim.state.a_prev + mdsim.state.a)

		if n % P.rescale_velocity_interval == 0 && P.ensemble_type == :NVT
			rescale_velocity!(mdsim)
		end

		compute_energies!(mdsim)
		update_history!(mdsim, n + 1)
	end
end


function formulation(
    i::MinimumBudget{SpanningTreeInstance{Int, Maximise}, Int},
    f::DefaultLinearFormulation
)
    m, x = formulation(i.instance, f)
    @constraint(
        m,
        c,
        sum((x[e] + x[reverse(e)]) * i.weights[e] for e in keys(i.weights)) >=
        0
    )
    return m, x, c
end

function solve(
    i::MinimumBudget{SpanningTreeInstance{Int, Maximise}, Int},
    f::DefaultLinearFormulation
)
    model, x, c = formulation(i, f)

    budgets = if i.compute_all_values == false
        [i.min_budget]
    else
        0:(i.min_budget)
    end

    S = Dict{Int, Vector{Edge{Int}}}()

    for budget in budgets
        set_normalized_rhs(c, budget)
        optimize!(model)

        if termination_status(model) == MOI.OPTIMAL
            S[budget] = _extract_lp_solution(i, x)
        else
            S[budget] = Int[-1]
        end
    end

    return BudgetedSpanningTreeDynamicProgrammingSolution(
        i,
        S[i.min_budget],
        S,
    )
end


# A quadratic segment in passes through three points: 𝐱₁, 𝐱₂, and 𝐱₃.
# The assumed relation of the points may be seen in the diagram below.
#                 ___𝐱₃___
#            ____/        \____
#        ___/                  \
#     __/                       𝐱₂
#   _/
#  /
# 𝐱₁
#
# NOTE: 𝐱₃ is not necessarily the midpoint, or even between 𝐱₁ and 𝐱₂.
# q(r) = (2r-1)(r-1)𝐱₁ + r(2r-1)𝐱₂ + 4r(1-r)𝐱₃
# See The Visualization Toolkit: An Object-Oriented Approach to 3D Graphics, 4th Edition
# Chapter 8, Advanced Data Representation, in the interpolation functions section
struct QuadraticSegment{F <: AbstractFloat} <: Edge{F}
    points::SVector{3, Point{F}}
end

# # Constructors
# # -------------------------------------------------------------------------------------------------
# QuadraticSegment(p₁::Point, p₂::Point, p₃::Point) = QuadraticSegment(SVector(p₁, p₂, p₃))
#
# # Methods
# # -------------------------------------------------------------------------------------------------
# # Interpolation
# # q(0) = q[1], q(1) = q[2], q(1//2) = q[3]
# function (q::QuadraticSegment{F})(r::Real) where {F <: AbstractFloat}
#     # See The Visualization Toolkit: An Object-Oriented Approach to 3D Graphics, 4th Edition
#     # Chapter 8, Advanced Data Representation, in the interpolation functions section
#     rₜ = F(r)
#     return (2rₜ-1)*( rₜ-1)q[1] + rₜ*(2rₜ-1)q[2] + 4rₜ*( 1-rₜ)q[3]
# end
#
# arclength(q::QuadraticSegment) = arclength(q, Val(15))
# function arclength(q::QuadraticSegment{F}, ::Val{N}) where {F <: AbstractFloat, N}
#     # Numerical integration is used.
#     # (Gauss-Legengre quadrature)
#     #     1                  N
#     # L = ∫ ||∇ q⃗(r)||dr  ≈  ∑ wᵢ||∇ q⃗(rᵢ)||
#     #     0                 i=1
#     #
#     w, r = gauss_legendre_quadrature(F, Val(N))
#     return sum(@. w * norm(∇(q, r)))
# end
#
# # Find the axis-aligned bounding box of the segment.
# function boundingbox(q::QuadraticSegment)
#     # Find the r coordinates where ∂x/∂r = 0, ∂y/∂r = 0
#     # We know ∇ q, so we can directly compute these values
#     r_x = (3q[1].x + q[2].x - 4q[3].x)/(4(q[1].x + q[2].x - 2q[3].x))
#     if 0 < r_x < 1
#         x_extreme = (2r_x-1)*(r_x-1)q[1].x + r_x*(2r_x-1)q[2].x + 4r_x*(1-r_x)q[3].x
#         xmin = min(q[1].x, q[2].x, x_extreme)
#         xmax = max(q[1].x, q[2].x, x_extreme)
#     else
#         xmin = min(q[1].x, q[2].x)
#         xmax = max(q[1].x, q[2].x)
#     end
#
#     r_y = (3q[1].y + q[2].y - 4q[3].y)/(4(q[1].y + q[2].y - 2q[3].y))
#     if 0 < r_y < 1
#         y_extreme = (2r_y-1)*(r_y-1)q[1].y + r_y*(2r_y-1)q[2].y + 4r_y*(1-r_y)q[3].y
#         ymin = min(q[1].y, q[2].y, y_extreme)
#         ymax = max(q[1].y, q[2].y, y_extreme)
#     else
#         ymin = min(q[1].y, q[2].y)
#         ymax = max(q[1].y, q[2].y)
#     end
#     return Rectangle(Point(xmin, ymin), Point(xmax, ymax))
# end
#
# # Return the gradient of q, evalutated at r
# function gradient(q::QuadraticSegment{F}, r::Real) where {F <: AbstractFloat}
#     #∇q =
#     rₜ = F(r)
#     return (4rₜ - 3)*(q[1] - q[3]) + (4rₜ - 1)*(q[2] - q[3])
# end
#
# # Return if the point is left of the quadratic segment
# #   p    ^
# #   ^   /
# # v⃗ |  / u⃗
# #   | /
# #   o
# # function isleft(p::Point, q::QuadraticSegment)
# #     if isstraight(q) || p ∉  boundingbox(q)
# #         # We don't need to account for the curve if q is straight or p is outside
# #         # q's bounding box
# #         u⃗ = q[2] - q[1]
# #         v⃗ = p - q[1]
# #         return u⃗ × v⃗ > 0
# #     else
# #         # Get the nearest point on q to p.
# #         # Construct vectors from a point on q, close to p_near, to p_near and p.
# #         # Use the cross product of these vectors to determine if p isleft.
# #         r, p_near = nearest_point(p, q)
# #
# #         if r < 1e-6 || 1 < r # If r is small or beyond the valid range, just use q[2]
# #             u⃗ = q[2] - q[1]
# #             v⃗ = p - q[1]
# #         else # otherwise use a point on q, close to p_near
# #             q_base = q(0.95r)
# #             u⃗ = p_near - q_base
# #             v⃗ = p - q_base
# #         end
# #         return u⃗ × v⃗ > 0
# #     end
# # end
#
# # If the quadratic segment is effectively linear
# @inline function isstraight(q::QuadraticSegment)
#     # u⃗ × v⃗ = |u⃗||v⃗|sinθ
#     return abs((q[3] - q[1]) × (q[2] - q[1])) < 1e-8
# end
#
# # function intersect(l::LineSegment, q::QuadraticSegment)
# #     ϵ = parametric_coordinate_ϵ
# #     if isstraight(q) # Use line segment intersection.
# #         # See LineSegment for the math behind this.
# #         v⃗ = l[2] - l[1]
# #         u⃗ = q[2] - q[1]
# #         vxu = v⃗ × u⃗
# #         # Parallel or collinear lines, return.
# #         1e-8 < abs(vxu) || return (0x00000000, SVector(Point(), Point()))
# #         w⃗ = q[1] - l[1]
# #         # Delay division until r,s are verified
# #         if 0 <= vxu
# #             lowerbound = (-ϵ)vxu
# #             upperbound = (1 + ϵ)vxu
# #         else
# #             upperbound = (-ϵ)vxu
# #             lowerbound = (1 + ϵ)vxu
# #         end
# #         r_numerator = w⃗ × u⃗
# #         s_numerator = w⃗ × v⃗
# #         if (lowerbound ≤ r_numerator ≤ upperbound) && (lowerbound ≤ s_numerator ≤ upperbound)
# #             return (0x00000001, SVector(l(s_numerator/vxu), Point()))
# #         else
# #             return (0x00000000, SVector(Point(), Point()))
# #         end
# #     else
# #         # q(r) = (2r-1)(r-1)𝐱₁ + r(2r-1)𝐱₂ + 4r(1-r)𝐱₃
# #         # q(r) = 2r²(𝐱₁ + 𝐱₂ - 2𝐱₃) + r(-3𝐱₁ - 𝐱₂ + 4𝐱₃) + 𝐱₁
# #         # Let D⃗ = 2(𝐱₁ + 𝐱₂ - 2𝐱₃), E⃗ = (-3𝐱₁ - 𝐱₂ + 4𝐱₃), F⃗ = x₁
# #         # q(r) = r²D⃗ + rE⃗ + F⃗
# #         # l(s) = 𝐱₄ + sw⃗
# #         # If D⃗ × w⃗ ≠ 0
# #         #   𝐱₄ + sw⃗ = r²D⃗ + rE⃗ + F⃗
# #         #   sw⃗ = r²D⃗ + rE⃗ + (F⃗ - 𝐱₄)
# #         #   0 = r²(D⃗ × w⃗) + r(E⃗ × w⃗) + (F⃗ - 𝐱₄) × w⃗
# #         #   Let A = (D⃗ × w⃗), B = (E⃗ × w⃗), C = (F⃗ - 𝐱₄) × w⃗
# #         #   0 = Ar² + Br + C
# #         #   r = (-B - √(B²-4AC))/2A, -B + √(B²-4AC))/2A)
# #         #   s = ((q(r) - p₄)⋅w⃗/(w⃗ ⋅ w⃗)
# #         #   r is invalid if:
# #         #     1) A = 0
# #         #     2) B² < 4AC
# #         #     3) r < 0 or 1 < r   (Curve intersects, segment doesn't)
# #         #   s is invalid if:
# #         #     1) s < 0 or 1 < s   (Line intersects, segment doesn't)
# #         # If D⃗ × w⃗ = 0, there is only one intersection and the equation reduces to line
# #         # intersection.
# #         npoints = 0x00000000
# #         p₁ = Point()
# #         p₂ = Point()
# #         D⃗ = 2(q[1] +  q[2] - 2q[3])
# #         E⃗ =  4q[3] - 3q[1] -  q[2]
# #         w⃗ = l[2] - l[1]
# #         A = D⃗ × w⃗
# #         B = E⃗ × w⃗
# #         C = (q[1] - l[1]) × w⃗
# #         w = w⃗ ⋅ w⃗
# #         if abs(A) < 1e-8
# #             # Line intersection
# #             # Can B = 0 if A = 0 for non-trivial 𝐱?
# #             r = -C/B
# #             (-ϵ ≤ r ≤ 1 + ϵ) || return 0x00000000, SVector(p₁, p₂)
# #             p₁ = q(r)
# #             s = (p₁ - l[1]) ⋅ w⃗/w
# #             if (-ϵ ≤ s ≤ 1 + ϵ)
# #                 npoints = 0x00000001
# #             end
# #         elseif B^2 ≥ 4A*C
# #             # Quadratic intersection
# #             # The compiler seem seems to catch the √(B^2 - 4A*C), for common subexpression
# #             # elimination, so leaving for readability
# #             r₁ = (-B - √(B^2 - 4A*C))/2A
# #             r₂ = (-B + √(B^2 - 4A*C))/2A
# #             if (-ϵ ≤ r₁ ≤ 1 + ϵ)
# #                 p = q(r₁)
# #                 s₁ = (p - l[1]) ⋅ w⃗/w
# #                 if (-ϵ ≤ s₁ ≤ 1 + ϵ)
# #                     p₁ = p
# #                     npoints += 0x00000001
# #                 end
# #             end
# #             if (-ϵ ≤ r₂ ≤ 1 + ϵ)
# #                 p = q(r₂)
# #                 s₂ = (p - l[1]) ⋅ w⃗/w
# #                 if (-ϵ ≤ s₂ ≤ 1 + ϵ)
# #                     p₂ = p
# #                     npoints += 0x00000001
# #                 end
# #             end
# #             if npoints === 0x00000001 && p₁ === Point()
# #                 p₁ = p₂
# #             end
# #         end
# #         return npoints, SVector(p₁, p₂)
# #     end
# # end
#
# # Return the Laplacian of q, evalutated at r
# function laplacian(q::QuadraticSegment, r::Real)
#     return 4(q[1] + q[2] - 2q[3])
# end
#
# nearest_point(p::Point, q::QuadraticSegment) = nearest_point(p, q, 15)
# # Return the closest point on the curve to point p and the value of r such that q(r) = p_nearest
# # Uses at most N iterations of Newton-Raphson
# function nearest_point(p::Point, q::QuadraticSegment, N::Int64)
#     r = 0.5
#     Δr = 0.0
#     for i = 1:N
#         err = p - q(r)
#         grad = ∇(q, r)
#         if abs(grad[1]) > abs(grad[2])
#             Δr = err[1]/grad[1]
#         else
#             Δr = err[2]/grad[2]
#         end
#         r += Δr
#         if abs(Δr) < 1e-7
#             break
#         end
#     end
#     return r, q(r)
# end
#
# #function newton_raphson(f::Function, J⁻¹::Function, x)
# #    xₙ₊₁ = xₙ - J⁻¹(xₙ)f
# #end
#
#
#
#
#
#
#
#
#
#
#
#
# # Plot
# # -------------------------------------------------------------------------------------------------
# if enable_visualization
#     function convert_arguments(LS::Type{<:LineSegments}, q::QuadraticSegment)
#         rr = LinRange(0, 1, 15)
#         points = q.(rr)
#         coords = reduce(vcat, [[points[i], points[i+1]] for i = 1:length(points)-1])
#         return convert_arguments(LS, coords)
#     end
#
#     function convert_arguments(LS::Type{<:LineSegments}, Q::Vector{<:QuadraticSegment})
#         point_sets = [convert_arguments(LS, q) for q in Q]
#         return convert_arguments(LS, reduce(vcat, [pset[1] for pset in point_sets]))
#     end
# end


"""
    q1_q0

The manufactured-solution colliding flow example from Elman et al 2014. The
linear quadrilaterals for the velocity and discontinuous
pressure.

The formulation is the one derived in Reddy, Introduction to the finite element
method, 1993. Page 486 ff.
"""
module q1_q0

using LinearAlgebra
using StaticArrays
using MeshCore
using MeshCore: retrieve, nrelations, nentities, ir_identity, attribute
using MeshSteward: Q4block
using MeshSteward: Mesh, attach!, baseincrel, boundary, increl
using MeshSteward: vselect, geometry, summary, transform
using MeshSteward: vtkwrite
using Elfel.RefShapes: manifdim, manifdimv
using Elfel.FElements: FEL2_Q4, FEH1_Q4, refshape, Jacobian
using Elfel.FESpaces: FESpace, ndofs, setebc!, nunknowns, doftype
using Elfel.FESpaces: numberfreedofs!, numberdatadofs!, numberdofs!
using Elfel.FESpaces: scattersysvec!, makeattribute, gathersysvec!, edofcompnt
using Elfel.FESpaces: highestfreedofnum, highestdatadofnum
using Elfel.FEIterators: FEIterator, ndofsperel, elnodes, eldofs
using Elfel.FEIterators: jacjac
using Elfel.QPIterators: QPIterator, bfun, bfungrad, weight
using Elfel.Assemblers: SysmatAssemblerSparse, start!, finish!, assemble!
using Elfel.Assemblers: SysvecAssembler
using Elfel.LocalAssemblers: LocalMatrixAssembler, LocalVectorAssembler, init!
using UnicodePlots

mu = 1.0 # dynamic viscosity
A = 1.0 # length of the side of the square
trueux = (x, y) -> 20 * x * y ^ 3
trueuy = (x, y) -> 5 * x ^ 4 - 5 * y ^ 4
truep = (x, y) -> 60 * x ^ 2 * y - 20 * y ^ 3

function genmesh(N)
    # This mesh will be both for the velocities and for the pressure
    mesh = Mesh()
    attach!(mesh, Q4block(2 * A, 2 * A, N, N), "velocity+pressure")
    ir = baseincrel(mesh)
    transform(ir, x -> x .- A)
    eidir = ir_identity(ir)
    attach!(mesh, eidir)
    return mesh
end

function assembleK(uxfesp, uyfesp, pfesp, tndof, mu)
    function integrateK!(ass, elits, qpits, mu)
        uxnedof, uynedof, pnedof = ndofsperel.(elits)
        kuxux = LocalMatrixAssembler(uxnedof, uxnedof, 0.0)
        kuyuy = LocalMatrixAssembler(uynedof, uynedof, 0.0)
        kuxuy = LocalMatrixAssembler(uxnedof, uynedof, 0.0)
        kuxp = LocalMatrixAssembler(uxnedof, pnedof, 0.0)
        kuyp = LocalMatrixAssembler(uynedof, pnedof, 0.0)
        for el in zip(elits...)
            uxel, uyel, pel = el
            init!(kuxux, eldofs(uxel), eldofs(uxel))
            init!(kuyuy, eldofs(uyel), eldofs(uyel))
            init!(kuxuy, eldofs(uxel), eldofs(uyel))
            init!(kuxp, eldofs(uxel), eldofs(pel))
            init!(kuyp, eldofs(uyel), eldofs(pel))
            for qp in zip(qpits...)
                uxqp, uyqp, pqp = qp
                Jac, J = jacjac(uxel, uxqp) # L2 pressure: must integrate over the velocity mesh
                JxW = J * weight(pqp)
                gradNp = bfungrad(pqp, Jac)
                gradNux = bfungrad(uxqp, Jac)
                gradNuy = bfungrad(uyqp, Jac)
                Np = bfun(pqp)
                for j in 1:uxnedof, i in 1:uxnedof
                    kuxux[i, j] += (mu * JxW) * (2 * gradNux[i][1] * gradNux[j][1] + gradNux[i][2] * gradNux[j][2])
                end
                for j in 1:uynedof, i in 1:uynedof
                    kuyuy[i, j] += (mu * JxW) * (gradNuy[i][1] * gradNuy[j][1] + 2 * gradNuy[i][2] * gradNuy[j][2])
                end
                for j in 1:uynedof, i in 1:uxnedof
                    kuxuy[i, j] += (mu * JxW) * (gradNux[i][1] * gradNuy[j][2])
                end
                for j in 1:pnedof, i in 1:uxnedof
                    kuxp[i, j] += (-JxW) * (gradNux[i][1] * Np[j])
                end
                for j in 1:pnedof, i in 1:uynedof
                    kuyp[i, j] += (-JxW) * (gradNuy[i][2] * Np[j])
                end
            end
            assemble!(ass, kuxux)
            assemble!(ass, kuxuy)
            assemble!(ass, transpose(kuxuy))
            assemble!(ass, kuyuy)
            assemble!(ass, kuxp)
            assemble!(ass, transpose(kuxp))
            assemble!(ass, kuyp)
            assemble!(ass, transpose(kuyp))
        end
        return ass
    end

    elits = (FEIterator(uxfesp), FEIterator(uyfesp), FEIterator(pfesp))
    qargs = (kind = :default, order = 2,)
    qpits = (QPIterator(uxfesp, qargs), QPIterator(uyfesp, qargs), QPIterator(pfesp, qargs))
    ass = SysmatAssemblerSparse(0.0)
    start!(ass, tndof, tndof)
    @time integrateK!(ass, elits, qpits, mu)
    return finish!(ass)
end

function solve!(U, K, F, nu)
    @time KT = K * U
    @time U[1:nu] = K[1:nu, 1:nu] \ (F[1:nu] - KT[1:nu])
end

centroid(p, c) = begin
    (p[c[1]] + p[c[2]] + p[c[3]] + p[c[4]]) / 4
end

function evaluate_error(uxfesp, uyfesp, pfesp)
    geom = geometry(pfesp.mesh)
    bir = baseincrel(pfesp.mesh)
    points = MeshCore.attribute(bir.right, "geom")
    centroids = [centroid(points, retrieve(bir, i)) for i in 1:nrelations(bir)]
    p = attribute(increl(pfesp.mesh, (2, 2)).right, "p")
    pt = [truep(centroids[i]...) for i in 1:length(centroids)]
    return norm(p - pt) / norm(pt)
end

function run(N)
    mesh = genmesh(N)
    # Velocity spaces
    uxfesp = FESpace(Float64, mesh, FEH1_Q4(), 1)
    uyfesp = FESpace(Float64, mesh, FEH1_Q4(), 1)
    locs = geometry(mesh)
    inflate = A / N / 100
    # The entire boundary
    boxes = [[-A A -A -A], [-A -A -A A], [A A -A A], [-A A A A]]
    for box in boxes
        vl = vselect(locs; box = box, inflate = inflate)
        for i in vl
            setebc!(uxfesp, 0, i, 1, trueux(locs[i]...))
            setebc!(uyfesp, 0, i, 1, trueuy(locs[i]...))
        end
    end
    # Pressure space
    pfesp = FESpace(Float64, mesh, FEL2_Q4(), 1)
    atcenter = vselect(geometry(mesh); nearestto = [0.0, 0.0])
    setebc!(pfesp, 2, atcenter[1], 1, 0.0)
    # Number the degrees of freedom
    numberdofs!([uxfesp, uyfesp, pfesp])
    @show tndof = ndofs(uxfesp) + ndofs(uyfesp) + ndofs(pfesp)
    @show tnunk = nunknowns(uxfesp) + nunknowns(uyfesp) + nunknowns(pfesp)
    # Assemble the coefficient matrix
    K = assembleK(uxfesp, uyfesp, pfesp, tndof, mu)
    # p = spy(K, canvas = DotCanvas)
    # display(p)
    # Solve the system
    U = fill(0.0, tndof)
    gathersysvec!(U, [uxfesp, uyfesp, pfesp])
    F = fill(0.0, tndof)
    solve!(U, K, F, tnunk)
    scattersysvec!([uxfesp, uyfesp, pfesp], U)
    # Postprocessing
    makeattribute(pfesp, "p", 1)
    makeattribute(uxfesp, "ux", 1)
    makeattribute(uyfesp, "uy", 1)
    @show evaluate_error(uxfesp, uyfesp, pfesp)
    vtkwrite("q1_q0-p", baseincrel(mesh), [(name = "p",), ])
    vtkwrite("q1_q0-v", baseincrel(mesh), [(name = "ux",), (name = "uy",)])
    true
end

end

let
    N = 4
    for loop in 1:5
        q1_q0.run(N)
        N = N * 2
    end
end


module RandomMatrices
importall Distributions
using Combinatorics

import Base.rand

#If the GNU Scientific Library is present, turn on additional functionality.
_HAVE_GSL = try
  using GSL
  true
catch
   false
end

export bidrand,    #Generate random bidiagonal matrix
       tridrand,   #Generate random tridiagonal matrix
       sprand,     #Generate random sparse matrix
       eigvalrand, #Generate random set of eigenvalues
       eigvaljpdf  #Eigenvalue joint probability density

typealias Dim2 (Int, Int) #Dimensions of a rectangular matrix

#Classical Gaussian matrix ensembles
include("GaussianEnsembles.jl")

# Classical univariate distributions
include("densities/Semicircle.jl")
include("densities/TracyWidom.jl")

# Ginibre
include("Ginibre.jl")

#Generating matrices of Haar measure
include("Haar.jl")
include("HaarMeasure.jl")

#Fast histogrammer for matrix eigenvalues - hist_eig
include("FastHistogram.jl")

#Formal power series
include("FormalPowerSeries.jl")

#Statistical tests based on random matrix theory
include("StatisticalTests.jl")

#Stochastic processes
include("StochasticProcess.jl")

#Invariant ensembles
if filemode(Pkg.dir("ApproxFun")) != 0
    include("InvariantEnsembles.jl")
end

end


# An eigenproblem described in the LaTeX file eigen.tex.

using FinElt
using FinElt.PlanarPoisson
using Arpack: eigs

const h = 0.05
run(`gmsh -2 -clmax $h -o roundL.msh -format msh22 roundL.geo`)

mesh = read_msh_file("roundL.msh")
essential_bc = [ "Dirichlet" ]
ep = EigenProblem(mesh, essential_bc)
add_bilin_form!(ep, "Omega", grad_dot_grad!,   :LHS)
add_bilin_form!(ep, "Omega", func_times_func!, :RHS)

A, B = assembled_eigenproblem_matrices(ep)

d, vfree, nconv, niter, nmult, resid = eigs(A, B, which=:SM)

v = complete_soln(vfree, ep)

write_pos_file("roundL.pos") do fid
    for k = 1:size(v,2)
        save_warp_nodal_scalar_field(v[:,k], "v$k", fid)
    end
end

run(`gmsh eigen.script`)


###################################################
## network/solvers/mip.jl
## MIP to select minimums
###################################################

"""
    `mipTimes` : optimize travel times to minimize L1 error from data with given paths
"""
function mipTimes(s::IterativeState; args...)
    g = s.data.network.graph
    paths = s.paths
    tripData = s.trips

    # Big M = max time of paths -- right now: 5km/h speed
    minTimes = uniformTimes(s.data.network, 5.)
    M = Float64[maximum(pathTime(minTimes,p) for p in l) for l in paths]

    #Create the model (will be changed to avoid hard-coded parameters)
    # !BarConvTol needs to be changed
    m = Model(solver = GurobiSolver(TimeLimit=10000, Method=3, BarConvTol=1e-6; args...))

    # DECISION VARIABLES
    # Road times
    @variable(m, t[i=vertices(g), j=outneighbors(g,i)] >= s.data.minTimes[i,j])
    # Absolute difference between tripData times and computed times
    @variable(m, epsilon[d=eachindex(tripData)] >= 0)
    # Time of the trips
    @variable(m, T[d=eachindex(tripData)] >= 0)
    # Integer variables to decide which path is the shortest
    @variable(m, minP[d=eachindex(tripData),p=1:length(paths[d])], Bin)

    # OBJECTIVE
    @objective(m, Min, sum(tripData[d].weight/tripData[d].time * epsilon[d] for d in eachindex(tripData)))

    # CONSTRAINTS
    # absolute values contraints (define epsilon), equal to time of first path
    @constraint(m, epsLower[d=eachindex(tripData)], T[d] - tripData[d].time >= - epsilon[d])
    @constraint(m, epsUpper[d=eachindex(tripData)], T[d] - tripData[d].time <=   epsilon[d])

    # inequality constraints
    @constraint(m, inequalityPath[d=eachindex(tripData), p=eachindex(paths[d])],
        sum(paths[d][p][edge] * t[src(edge), dst(edge)] for edge in keys(paths[d][p])) >= T[d] )

    # "minimum equality" contraints
    @constraint(m, equalityPath[d=eachindex(tripData), p=eachindex(paths[d])],
        sum(paths[d][p][edge] * t[src(edge), dst(edge)] for edge in keys(paths[d][p])) - M[d]*(1 - minP[d,p]) <= T[d])

    # integer constraints
    @constraint(m, equalityPath[d=eachindex(tripData)], sum(minP[d,p] for p in eachindex(paths[d])) == 1)

    # SOLVE LP
    status = solve(m)
    times = getvalue(t)

    # Export result as sparse matrix
    result = spzeros(Float64, nv(g), nv(g))
    for i in vertices(g), j in outneighbors(g,i)
        result[i,j] = times[i,j]
    end

    return result
end


using SweepContractor

L=16; d=2

LTN = LabelledTensorNetwork{Tuple{Int,Int}}()
for i∈1:L, j∈1:L
    adj=Tuple{Int,Int}[];
    i>1 && push!(adj,(i-1,j))
    j>1 && push!(adj,(i,j-1))
    i<L && push!(adj,(i+1,j))
    j<L && push!(adj,(i,j+1))
    LTN[i,j] = Tensor(adj, randn(d*ones(Int,length(adj))...), i, j)
end

println("example_2d_open")
for χ ∈ d.^(4:9)
    sweep = sweep_contract(LTN, χ, 2*χ; fast=true)
    println("χ=$χ:\t", ldexp(sweep...))
end