##### OVERRIDES FOR EFFICIENCY / CORRECTNESS

function add_vertices!(g::AbstractSimpleWeightedGraph, n::Integer)
    T = eltype(g)
    U = weighttype(g)
    (nv(g) + one(T) <= nv(g)) && return false       # test for overflow
    emptycols = spzeros(U, nv(g) + n, n)
    g.weights = hcat(g.weights, emptycols[1:nv(g), :])
    g.weights = vcat(g.weights, emptycols')
    return true
end

function adjacency_matrix(g::AbstractSimpleWeightedGraph, T::DataType=Int; dir::Symbol=:out)
    if dir == :out
        return SparseMatrixCSC(T.(LinearAlgebra.fillstored!(copy(g.weights), 1))')
    else
        return T.(LinearAlgebra.fillstored!(copy(g.weights), 1))
    end
end

function pagerank(g::SimpleWeightedDiGraph, α=0.85, n=100::Integer, ϵ=1.0e-6)
    A = weights(g)
    S = vec(sum(A, dims=1))
    S = 1 ./ S
    S[findall(S .== Inf)] .= 0.0
    M = A'  # need a separate line due to bug #17456 in julia
    # scaling the adjmat to stochastic adjacency matrix
    M = (Diagonal(S) * M)'
    N = Int(nv(g))
    # solution vector
    x = fill(1.0 / N, N)
    # personalization vector
    p = fill(1.0 / N, N)
    # temporary to hold the results of SpMV
    y = zeros(Float64, N)
    # adjustment for leaf nodes in digraph
    dangling_weights = p
    is_dangling = findall(S .== 0)
    # save some flops by precomputing this
    pscaled = (1 .- α) .* p
    for _ in 1:n
        xlast = x
        # in place SpMV to conserve memory
        mul!(y, M, x)
        # using broadcast to avoid temporaries
        x = α .* (y .+ sum(x[is_dangling]) .* dangling_weights) .+ pscaled
        # l1 change in solution convergence criterion
        err = sum(abs, (x .- xlast))
        if (err < N * ϵ)
            return x
        end
    end
    error("Pagerank did not converge after $n iterations.")
end

savegraph(fn::AbstractString, g::AbstractSimpleWeightedGraph, gname::AbstractString="graph"; compress=true) =
    savegraph(fn, g, gname, SWGFormat(), compress=compress)

savegraph(fn::AbstractString, d::Dict{T, U}; compress=true) where T <: AbstractString where U <: AbstractSimpleWeightedGraph =
    savegraph(fn, d, SWGFormat(), compress=compress)

# Connected Components on a Sparse Matrix

function _cc(g::SimpleWeightedGraph{T,U}) where T where U
    a = weights(g)
    comp = 0
    n = size(a, 1)
    marks = zeros(T, n)
    queue = Vector{T}()
    for i = 1:n
        if marks[i] == 0
            comp += 1
            push!(queue, i)
            while !isempty(queue)
                v = pop!(queue)
                marks[v] = comp
                for index in nzrange(a,v)
                    n = a.rowval[index]
                    if marks[n] == 0
                        push!(queue, n)
                    end
                end
            end
        end
    end
    marks, comp
end

function connected_components(g::SimpleWeightedGraph{T,U}) where T where U
    marks, num_cc = _cc(g)
    cc = [Vector{T}() for i = 1:num_cc]
    for (i,v) in enumerate(marks)
        push!(cc[v], i)
    end
    cc
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.10.3
#   kernelspec:
#     display_name: Julia 1.6.3
#     language: julia
#     name: julia-1.6
# ---

# %%
using AbstractAlgebra
AbstractAlgebra.PrettyPrinting.set_html_as_latex(true)

eqstr(x) = "\\ \\text{is " * string(x) * "}"
eqstr(x::SetElem) = "=" * sprint(x) do io, x; show(io, "text/latex", x) end

function dispeq(lhs, rhs)
    latexstr = "\$\\displaystyle $lhs " * eqstr(rhs) * "\$"
    display("text/html", latexstr)
end

function euclidean(f, g, stop = 1)
    while true
        r = mod(f, g)
        iszero(r) && return g
        degree(r) ≤ stop && return r
        f, g = g, r
    end
end

function rootdeg1(f)
    C = coefficients(f) |> collect
    -first(C)//last(C)
end

function commonroot(f, g)
    m = euclidean(f, g, 1)
    rootdeg1(m)
end

function revpoly(z, x, f)
    m = degree(f)
    C = coefficients(f)
    sum(a*z^(i-1)*x^(m-i+1) for (i, a) in enumerate(C))
end

function val(h::FracElem, θ, L)
    n = L(numerator(h)(θ))
    d = L(denominator(h)(θ))
    n//d
end

function testval(h::FracElem, θ, β, L)
    n = L(numerator(h)(θ))
    d = L(denominator(h)(θ))
    n == β*d
end

function calcallresults(z, x, f, g)
    f_plus = (-1)^degree(f)*f(z - x)
    F_plus = resultant(f_plus, g)
    β_plus = commonroot(g, f_plus)

    f_mult = (-1)^degree(f)*revpoly(z, x, f)
    F_mult = resultant(f_mult, g)
    β_mult = commonroot(g, f_mult)

    K = FractionField(base_ring(z))
    B1, α = K["α"]
    L1 = ResidueField(B1, numerator(f(z))(α))
    B2, β = L1["β"]
    L = ResidueField(B2, numerator(g(z))(β))

    val_F_plus = val(F_plus, α + β, L)
    test_β_plus = testval(β_plus, α + β, β, L)
    val_F_mult = val(F_mult, α * β, L)
    test_β_mult = testval(β_mult, α * β, β, L)

    F_plus, β_plus, F_mult, β_mult, val_F_plus, test_β_plus, val_F_mult, test_β_mult
end

function dispallresults(z, x, f, g)
    F_plus, β_plus, F_mult, β_mult, val_F_plus, test_β_plus, val_F_mult, test_β_mult =
        @time calcallresults(z, x, f, g)
    flush(stdout)
    dispeq("F_\\alpha(x)", f)
    dispeq("F_\\beta(x)", g)
    dispeq("R_{\\alpha + \\beta}(z)", F_plus)
    dispeq("R_{\\alpha + \\beta}(\\alpha + \\beta)", val_F_plus)
    dispeq("\\beta_\\mathrm{plus}(z)", β_plus)
    dispeq("\\beta_\\mathrm{plus}(\\alpha + \\beta) = \\beta", test_β_plus)
    dispeq("R_{\\alpha\\beta}(z)", F_mult)
    dispeq("R_{\\alpha\\beta}(\\alpha\\beta)", val_F_mult)
    dispeq("\\beta_\\mathrm{mult}(z)", β_mult)
    dispeq("\\beta_\\mathrm{mult}(\\alpha\\beta) = \\beta", test_β_mult)
end

# %%
R, (a, b, c, p, q, r, s) = ZZ["a", "b", "c", "p", "q", "r", "s"]
K = FractionField(R)
Rz, z = R["z"]
Kz = FractionField(Rz)
Rx, x = Kz["x"]

# %%
f = x^2 - a
g = x^2 - p
dispallresults(z, x, f, g)

# %%
f = x^2 - a
g = x^3 - p
dispallresults(z, x, f, g)

# %%
f = x^2 - a
g = x^4 - p
dispallresults(z, x, f, g)

# %%
f = x^2 - a
g = x^5 - p
dispallresults(z, x, f, g)

# %%
f = x^2 - a
g = x^6 - p
dispallresults(z, x, f, g)

# %%
f = x^2 - a
g = x^7 - p
dispallresults(z, x, f, g)

# %%
f = x^3 - a
g = x^4 - p
dispallresults(z, x, f, g)

# %%
f = x^3 - a
g = x^5 - p
dispallresults(z, x, f, g)

# %%
f = x^3 - a
g = x^6 - p
dispallresults(z, x, f, g)

# %%
f = x^3 - a
g = x^7 - p
dispallresults(z, x, f, g)

# %%
f = x^4 - a
g = x^4 - p
dispallresults(z, x, f, g)

# %%
f = x^4 - a
g = x^5 - p
dispallresults(z, x, f, g)

# %%
f = x^5 - a
g = x^5 - p
dispallresults(z, x, f, g)

# %%
f = x^2 - a
g = x^3 + p*x + q
dispallresults(z, x, f, g)

# %%
f = x^2 - a
g = x^4 + p*x^2 + q*x + r
dispallresults(z, x, f, g)

# %%
f = x^2 - a
g = x^5 + p*x^3 + q*x^2 + r*x + s
dispallresults(z, x, f, g)

# %%
f = a*x^2 + b*x + c
g = p*x^2 + q*x + r
dispallresults(z, x, f, g)

# %%
f = x^2 + b*x + c
g = x^2 + q*x + r
dispallresults(z, x, f, g)

# %%
f = x^2 + b*x + c
g = x^3 + p*x^2 + q*x + r
dispallresults(z, x, f, g)

# %%
f = x^2 + a*x + b
g = x^4 + p*x^3 + q*x^2 + r*x + s
dispallresults(z, x, f, g)

# %%
f = x^2 + a*x + b
g = x^5 + p*x^3 + q*x^2 + r*x + s
dispallresults(z, x, f, g)

# %%
f = x^2 + a*x + b
g = x^6 + p*x^3 + q*x^2 + r*x + s
dispallresults(z, x, f, g)

# %%
f = x^2 + a*x + b
g = x^7 + p*x^3 + q*x^2 + r*x + s
dispallresults(z, x, f, g)

# %%
f = x^3 - a
g = x^3 + p*x + q
dispallresults(z, x, f, g)

# %%
f = x^3 - a
g = x^4 + p*x^2 + q*x + r
dispallresults(z, x, f, g)

# %%
f = x^3 - a
g = x^5 + p*x^3 + q*x^2 + r*x + s
dispallresults(z, x, f, g)

# %%
f = x^3 + a*x + b
g = x^3 + p*x + q
dispallresults(z, x, f, g)

# %%
f = x^3 + a*x + b
g = x^4 + p*x^2 + q*x + r
dispallresults(z, x, f, g)

# %%


@testset "Neighborhoods" begin
  @testset "BallNeighborhood" begin
    # Euclidean metric
    n = BallNeighborhood{1}(0.5)
    @test isneighbor(n, [0.], [0.])
    @test !isneighbor(n, [0.], [1.])

    n = BallNeighborhood{2}(1.0)
    @test isneighbor(n, [0.,0.], [0.,0.])
    @test isneighbor(n, [0.,0.], [1.,0.])
    @test isneighbor(n, [0.,0.], [0.,1.])

    # Chebyshev metric
    n = BallNeighborhood{1}(0.5, Chebyshev())
    @test isneighbor(n, [0.], [0.])
    @test !isneighbor(n, [0.], [1.])

    for r in [1.,2.,3.,4.,5.]
      n = BallNeighborhood{2}(r, Chebyshev())
      for i in 0.0:1.0:r, j in 0.0:1.0:r
        @test isneighbor(n, [0.,0.], [i,j])
      end
    end

    n = BallNeighborhood{2}(1.0)
    @test sprint(show, n) == "2D BallNeighborhood(1.0, Euclidean(0.0))"
    @test sprint(show, MIME"text/plain"(), n) == "2D BallNeighborhood\n  radius: 1.0\n  metric: Euclidean(0.0)"
  end

  @testset "CylinderNeighborhood" begin
    g = RegularGrid{Float64}(3,3,3)
    n = CylinderNeighborhood{3}(1., 1.)
    for i in [1,2,4,10,11,13]
      @test isneighbor(n, coordinates(g, 1), coordinates(g, i))
    end
    n = CylinderNeighborhood{3}(2., 1.)
    for i in [1,2,3,4,5,7,10,11,12,13,14,16]
      @test isneighbor(n, coordinates(g, 1), coordinates(g, i))
    end
    n = CylinderNeighborhood{3}(1., 2.)
    for i in [1,2,4,10,11,13,19,20,22]
      @test isneighbor(n, coordinates(g, 1), coordinates(g, i))
    end

    g = RegularGrid{Float64}(100,100,100)
    n = CylinderNeighborhood{3}(1., 1.)
    for i in [1,2,101,10001,10002,10101]
      @test isneighbor(n, coordinates(g, 1), coordinates(g, i))
    end

    n = CylinderNeighborhood{3}(1., 1.)
    @test sprint(show, n) == "3D CylinderNeighborhood(1.0, 1.0)"
    @test sprint(show, MIME"text/plain"(), n) == "3D CylinderNeighborhood\n  radius: 1.0\n  height: 1.0"
  end
end


# TODO: allow for different shapes and mirror sizes

using Plots
using LinearAlgebra

pyplot()

mutable struct Ray
    # Direction vector
    l::Vector{Float64}

    # Position vector
    p::Vector{Float64}

    # Current Index of Refraction
    ior::Float64
end

mutable struct Lens
    # Lens position
    p::Vector{Float64}

    # Lens radius
    r::Float64

    # Lens function
    f
end

function constant(p)
    return 1.5
end

function mag_index(p)
    return cos(sum(p[:].^2))
end

function inverse(p)
    mag = sum(p[:].^2)
    return 1/(mag)
end

function lens_normal_at(ray, lens)
    n = normalize(ray.p .- lens.p)
    if dot(-n, ray.l) < 0
        n *= -1
    end

    return n
end

function inside_of(ray::Ray, lens)
    return inside_of(ray.p, lens)
end

function inside_of(pos, lens)
    x = lens.p[1] - pos[1]
    y = lens.p[2] - pos[2]
    if (x^2 + y^2 <= lens.r^2)
        return true
    else
        return false
    end
end

function refract!(ray, n)
    c = dot(-n, ray.l);
    d = 1.0 - ior^2 * (1.0 - c^2);

    if (d < 0.0)
        return zeros(2);
    end

    ray.l = ior * ray.l + (ior * c - sqrt(d)) * n;
end

# note: light moves as a particular speed with respect to the medium it is
#       moving through, so...
#       n_2*v = n_1*l + (n_1*cos(theta_1) - n_2*cos(theta_2))*n
#       Other approximations: ior = n_1/n_2, c = -n*l
function refract!(ray, lens, ior)
    n = lens_normal_at(ray, lens)
    c = dot(-n, ray.l);
    d = 1.0 - ior^2 * (1.0 - c^2);

    if (d < 0.0)
        return zeros(2);
    end

    ray.l = ior * ray.l + (ior * c - sqrt(d)) * n;
end

mutable struct Mirror
    # Normal vector
    n::Vector{Float64}

    # Position of mirror
    p::Vector{Float64}

    # Mirror size
    scale::Float64

    Mirror(in_n, in_p) = new(in_n, in_p, 2.5)
end

function is_behind(ray, mirror)
    if dot(ray.p.-mirror.p, mirror.n) >= 0
        return true
    else
        return false
    end
end

# note: for reflection, l_x -> l_x, but l_y -> -l_y
#       In this case, the y component of l = cos(theta)*n
#       so new vector: v = l + 2cos(theta)*n
function reflect!(ray, n)
    ray.l .-= 2*dot(ray.l, n).*n
end

function draw_circle(p, r, res)
    return [x .+ (p[1], p[2]) for x in Plots.partialcircle(0, 2pi, res, r)]
end

function plot_rays(positions, mirrors::Vector{Mirror},
                   lenses::Vector{Lens}, filename)
    plt = plot(background_color=:black, aspect_ratio=:equal)
    for i = 1:size(positions)[1]
        plot!(plt, (positions[i,:,1], positions[i,:,2]); label = "ray",
              linecolor=:white)
    end

    for mirror in mirrors
        dir = [-mirror.n[2], mirror.n[1]]
        extents = zeros(2,2)
        extents[1,:] .= mirror.p .- mirror.scale * dir
        extents[2,:] .= mirror.p .+ mirror.scale * dir
        plot!(plt, (extents[:,1], extents[:,2]), label = "mirror")
    end

    for lens in lenses
        circle = draw_circle(lens.p, lens.r, 100)
        plot!(circle; label="lens", linecolor=:lightblue)
    end

    savefig(plt, filename)
end

function step!(ray::Ray, dt)
    ray.p .+= .+ ray.l.*dt/ray.ior
end

function propagate!(rays::Vector{Ray}, mirrors::Vector{Mirror},
                    lenses::Vector{Lens}, positions, n, dt)

    ior = 1.0
    for i = 1:n
        for j = 1:length(rays)
            for mirror in mirrors
                if is_behind(rays[j], mirror)
                    reflect!(rays[j], mirror.n)
                end
            end
            for lens in lenses
                if inside_of(rays[j], lens) && i > 2
                   ior = rays[j].ior/lens.f(rays[j].p .- lens.p)
                   rays[j].ior = lens.f(rays[j].p .- lens.p)
                   refract!(rays[j], lens, ior)
                elseif !inside_of(rays[j], lens) && i > 2 &&
                       rays[j].ior != 1
                    ior = rays[j].ior
                    rays[j].ior = 1
                    refract!(rays[j], lens, ior)
                end
                if rays[j].l == zeros(2)
                    reflect!(rays[j], lens_normal_at(rays[j], lens))
                end
            end
            step!(rays[j], dt)
            positions[j, i,:] .= rays[j].p
        end
    end
end

function parallel_propagate(ray_num, n; filename="check.png", dt = 0.1)

    rays = [Ray([1, 0],
            [0.1, float(i)],
            1.0) for i = 1:ray_num]

    #lenses = [Lens([10, 6], 6, constant)]
    lenses = [Lens([10, 6], 6, constant)]
    #lenses = [Lens([-15, -10], 5, 1.5)]
    mirrors = [Mirror([0.0, -1.0],[5.0, 0.0])]
#=
    mirrors = [Mirror([0.0, 1.0],[5.0, 5.0]),
               Mirror([0.0, -1.0],[5.0, 0.0]),
               Mirror([1.0, 0.0],[0.0, 2.5]),
               Mirror([-1.0, 0.0],[10, 2.5])]
=#

    positions = zeros(ray_num, n, 2)

    propagate!(rays, mirrors, lenses, positions, n, dt)

    plot_rays(positions, mirrors, lenses, filename)

    return rays
end


using Stheno, Plots, Random
using Stheno: @model

###########################  Define and inspect our model  ###########################

#=

=#
@model function model()
    g1, g2 = x::Real->sin(x), x::Real->x
    w1, w2 = GP(EQ()), GP(EQ())
    f = g1 * w1 + g2 * w2
    y = f + GP(Noise(0.001))
    return w1, w2, f, y
end

# Sample from the prior from plotting and for conditioning.
rng, N, Nplot, S = MersenneTwister(123456), 250, 500, 100;
X, Xp = sort(rand(rng, N) * 10), range(-2.5, stop=12.5, length=Nplot);
w1, w2, f, y = model();
w1s, w2s, fs, ŷ = rand(rng, [w1(Xp), w2(Xp), f(Xp), y(X)]);

# Compute posterior distribution over f′.
w1′, w2′, f′, y′ = (w1, w2, f, y) | (y(X) ← ŷ);

# Sample from the posterior and write to file.
# w1′s, w2′s, f′s = rand(rng, [w1′(Xp), w2′(Xp), f′(Xp)], S);
w1′s, w2′s, y′s = rand(rng, [w1′(Xp), w2′(Xp), y′(Xp)], S);

# Get posterior mean and marginals f′ and y′ and write them for plotting.
μw1′, σw1′ = marginals(w1′(Xp));
μw2′, σw2′ = marginals(w2′(Xp));
# μf′, σf′ = marginals(f′(Xp));
μy′, σy′ = marginals(y′(Xp));

###########################  Plot results - USE ONLY Julia-0.6!  ###########################

posterior_plot = plot();

# Plot samples against which we're regressing.
plot!(posterior_plot, X, ŷ;
    markercolor=:red,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=4,
    markeralpha=0.7,
    label="");

# # Plot posterior over `f`.
# plot!(posterior_plot, Xp, μf′;
#     linecolor=:blue,
#     linewidth=2.0,
#     label="f");
# plot!(posterior_plot, Xp, [μf′ μf′];
#     linewidth=0.0,
#     fillrange=[μf′ .- 3 .* σf′, μf′ .+ 3 * σf′],
#     fillalpha=0.3,
#     fillcolor=:blue,
#     label="");
# plot!(posterior_plot, Xp, f′s;
#     linecolor=:blue,
#     linealpha=0.2,
#     label="");

# Plot posterior over `y`.
plot!(posterior_plot, Xp, μy′;
    linecolor=:blue,
    linewidth=2.0,
    label="y");
plot!(posterior_plot, Xp, [μy′ μy′];
    linewidth=0.0,
    fillrange=[μy′ .- 3 .* σy′, μy′ .+ 3 * σy′],
    fillalpha=0.3,
    fillcolor=:blue,
    label="");
plot!(posterior_plot, Xp, y′s;
    linecolor=:blue,
    linealpha=0.2,
    label="");

# Plot posterior over w1.
plot!(posterior_plot, Xp, μw1′;
    linecolor=:green,
    linewidth=2.0,
    label="w1");
plot!(posterior_plot, Xp, [μw1′ μw1′];
    linewidth=0.0,
    fillrange=[μw1′ .- 3 .* σw1′, μw1′ .+ 3 * σw1′],
    fillalpha=0.3,
    fillcolor=:green,
    label="");
plot!(posterior_plot, Xp, w1′s;
    linecolor=:green,
    linealpha=0.2,
    label="");

# Plot posterior over w2.
plot!(posterior_plot, Xp, μw2′;
    linecolor=:magenta,
    linewidth=2.0,
    label="w2");
plot!(posterior_plot, Xp, [μw2′ μw2′];
    linewidth=0.0,
    fillrange=[μw2′ .- 3 .* σw2′, μw2′ .+ 3 * σw2′],
    fillalpha=0.3,
    fillcolor=:magenta,
    label="");
plot!(posterior_plot, Xp, w2′s;
    linecolor=:magenta,
    linealpha=0.2,
    label="");

display(posterior_plot);


function genTikhonov(A,mr,p,Aufloesung)
lrows,lcols = size(A);
L = zeros(lcols,lcols);
for r = 1:lcols-1
    L[r,r] = -2;
    L[r,r+1] = 1;
    if r <= (Aufloesung^2-Aufloesung)
        L[r,r+Aufloesung] = 1;
    end
end
L[lcols,lcols] = -2;

ATA = transpose(A)*A;
LTL = transpose(L)*L;
GenTikVek = (ATA+p*LTL)\transpose(A)*mr;

return GenTikVek;
end


################################################################################
#
# Helmholtz-related linear algebra routines
#
################################################################################

export HelmholtzLinOp, computeNabc

function HelmholtzLinOp(h::NTuple{2, R}, m::Array{R, 2}, freq::R, nabc::NTuple{4, Int64})

    # Extension for ABC
    m_ext = extend(m, nabc)
    n_ext = size(m_ext)

    # Complex frequency for ABC
    abc_fact = 1f0
    ω = 2f0*pi*freq*(1f0.-abc_fact*1im*taper(size(m), nabc))

    # Laplacian
    Dxx = spdiagm(-1 => ones(R, n_ext[1]-1), 0 => -2f0*ones(R, n_ext[1]), 1 => ones(R, n_ext[1]-1))/h[1]^2
    Dxx[1, 1:2] = [1 -1]/h[1]^2 # Sommerfeld radiation condition
    Dxx[end, end-1:end] = [-1, 1]/h[1]^2 # Sommerfeld radiation condition
    Dzz = spdiagm(-1 => ones(R, n_ext[2]-1), 0 => -2f0*ones(R, n_ext[2]), 1 => ones(R, n_ext[2]-1))/h[2]^2
    Dzz[1, 1:2] = [1 -1]/h[2]^2 # Sommerfeld radiation condition
    Dzz[end, end-1:end] = [-1, 1]/h[2]^2 # Sommerfeld radiation condition
    Δ = kron(sparse(I, n_ext[2], n_ext[2]), Dxx)+kron(Dzz, sparse(I, n_ext[1], n_ext[1]))

    # Mass matrix
    w = ones(R, n_ext)
    w[[1, end], :] .= 0f0
    w[:, [1, end]] .= 0f0
    v = 1f0.-w
    v[:, [1, end]] = v[:, [1, end]]/h[2] # Sommerfeld radiation condition
    v[[1, end], :] = v[[1, end], :]/h[1] # Sommerfeld radiation condition
    # M = spdiagm(0 => vec(ω.^2f0.*w.*m_ext))+1im*spdiagm(0 => vec(ω.*v.*sqrt.(m_ext)))
    M = spdiagm(0 => vec(ω.^2f0.*w.*m_ext))+1im*spdiagm(0 => vec(ω.*v.*sqrt.(C.(m_ext))))

    # Assembly
    return -M-Δ

end

function taper(n::Tuple{Int64, Int64}, nabc::NTuple{4, Int64})

    taper_x = 1f0.-cat([range(-1f0, stop = 0f0, length = nabc[1])...].^2, zeros(R, n[1]), [range(0f0, stop = 1f0, length = nabc[2])...].^2, dims = 1)
    taper_z = 1f0.-cat([range(-1f0, stop = 0f0, length = nabc[3])...].^2, zeros(R, n[2]), [range(0f0, stop = 1f0, length = nabc[4])...].^2, dims = 1)
    return 1f0.-reshape(taper_x, :, 1)*reshape(taper_z, 1, :)

end

function computeNabc(h::NTuple{2, R}, m::Array{R, 2}, freq::R)

    return Int64.((round(1f0./(min(h...)*min(sqrt.(m[1, :])...)*freq)), round(1f0./(min(h...)*min(sqrt.(m[end, :])...)*freq)), round(1f0./(min(h...)*min(sqrt.(m[:, 1])...)*freq)), round(1f0./(min(h...)*min(sqrt.(m[:, end])...)*freq))))

end


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:2] <= 1)

                     @variable(m, -6.7172365869297215 <= q <= 3.8657529446129524)

                     add_NL_constraint(m, :(swish(0.10971879992509148 + -0.31290714990830937*$(x[1]) + 0.781487131205628*$(x[2])) + swish(-0.5708088154317625 + 0.5853291463969379*$(x[1]) + -0.9903406373931247*$(x[2])) + swish(0.5809119298414265 + 0.5024686961828384*$(x[1]) + -0.4096226237713658*$(x[2])) + swish(-0.6435921083509291 + 0.52943737688555*$(x[1]) + 0.2727617072630939*$(x[2])) + swish(-0.9019716271422107 + 0.4041768804112231*$(x[1]) + 0.5029634163532655*$(x[2])) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


# Function that updates partial force and energy vectors, used for parallel computation

function compute_uf_partials_LJ!(it :: Int64, f :: Forces,
                                 i :: Int64, j :: Int64,
                                 xj :: Float64, yj :: Float64, r :: Float64,
                                 input :: MDInput)
  r6 = r^6
  r12 = r6^2
  r7 = r6*r
  r13 = r12*r

  f.upartial[it] = f.upartial[it] + f.eps4*( f.sig12 / r12 - f.sig6 / r6 )

  drdxi = -xj/r
  drdyi = -yj/r

  fx = -f.eps4*( (f.sig1212/r13) - (f.sig66/r7) )*drdxi
  fy = -f.eps4*( (f.sig1212/r13) - (f.sig66/r7) )*drdyi

  f.fpartial[it][i,1] = f.fpartial[it][i,1] + fx
  f.fpartial[it][i,2] = f.fpartial[it][i,2] + fy
  f.fpartial[it][j,1] = f.fpartial[it][j,1] - fx
  f.fpartial[it][j,2] = f.fpartial[it][j,2] - fy

end


using JuMP, EAGO, Gurobi

m = Model(with_optimizer(EAGO.Optimizer, relaxed_optimizer = Gurobi.Optimizer(OutputFlag=0)))

# ----- Variables ----- #
x_Idx = Any[2, 3, 4, 5, 6, 7]
@variable(m, x[x_Idx])
JuMP.set_lower_bound(x[2], 1.0e-7)
JuMP.set_upper_bound(x[2], 0.08)
JuMP.set_lower_bound(x[3], 1.0e-7)
JuMP.set_upper_bound(x[3], 0.08)
JuMP.set_lower_bound(x[4], 1.0e-7)
JuMP.set_upper_bound(x[4], 0.3)
JuMP.set_lower_bound(x[5], 1.0e-7)
JuMP.set_upper_bound(x[5], 0.3)
JuMP.set_lower_bound(x[6], 1.0e-7)
JuMP.set_upper_bound(x[6], 0.62)
JuMP.set_lower_bound(x[7], 1.0e-7)
JuMP.set_upper_bound(x[7], 0.62)

# ----- Constraints ----- #
@constraint(m, e2, x[2]+x[3] == 0.08)
@constraint(m, e3, x[4]+x[5] == 0.3)
@constraint(m, e4, x[6]+x[7] == 0.62)

# ----- Objective ----- #
@NLobjective(m, Min, (x[2]*log(x[2]/(3*x[2]+6*x[4]+x[6]))+x[4]*log(x[4]/(3*x[2]+6*x[4]+x[6]))+x[6]*log(x[6]/(3*x[2]+6*x[4]+x[6]))-0.80323071133189*x[2]+1.79175946922805*x[4]+0.752006*x[6]+(3*x[2]+6*x[4]+1.6*x[6])*log(3*x[2]+6*x[4]+1.6*x[6])+2*x[2]*log(x[2]/(2.00000019368913*x[2]+4.64593*x[4]+0.480353*x[6]))+x[2]*log(x[2]/(1.00772874182154*x[2]+0.724703350369523*x[4]+0.947722362492017*x[6]))+6*x[4]*log(x[4]/(3.36359157977228*x[2]+6*x[4]+1.13841069150863*x[6]))+1.6*x[6]*log(x[6]/(1.6359356134845*x[2]+3.39220996773471*x[4]+1.6*x[6]))+x[3]*log(x[3]/(3*x[3]+6*x[5]+x[7]))+x[5]*log(x[5]/(3*x[3]+6*x[5]+x[7]))+x[7]*log(x[7]/(3*x[3]+6*x[5]+x[7]))-0.80323071133189*x[3]+1.79175946922805*x[5]+0.752006*x[7]+(3*x[3]+6*x[5]+1.6*x[7])*log(3*x[3]+6*x[5]+1.6*x[7])+2*x[3]*log(x[3]/(2.00000019368913*x[3]+4.64593*x[5]+0.480353*x[7]))+x[3]*log(x[3]/(1.00772874182154*x[3]+0.724703350369523*x[5]+0.947722362492017*x[7]))+6*x[5]*log(x[5]/(3.36359157977228*x[3]+6*x[5]+1.13841069150863*x[7]))+1.6*x[7]*log(x[7]/(1.6359356134845*x[3]+3.39220996773471*x[5]+1.6*x[7]))-3*x[2]*log(x[2])-6*x[4]*log(x[4])-1.6*x[6]*log(x[6])-3*x[3]*log(x[3])-6*x[5]*log(x[5])-1.6*x[7]*log(x[7])))

JuMP.optimize!(m)
run_time = backend(m).optimizer.model.optimizer._run_time
println("run time: $run_time")


using Romeo
using Plots

α = 0.5
β = 0.5
Δt = 0.01
N = 100

t = Δt:Δt:((N - 1) * Δt)
g_c = t -> Romeo.Fractional.ComplexOrder.kernel_fcn_co_integral(t, α, β)
g = Romeo.Operators.kernel_c2d(g_c, t)
g_i = t -> Romeo.Fractional.ComplexOrder.kernel_fcn_co_integral(t, α + 1, β)
gi = Romeo.Operators.kernel_from_integral(g_i, t)
u = ones(size(t))
y = Romeo.Operators.response(g, u)
yi = Romeo.Operators.response(gi, u)
ye = Romeo.Fractional.ComplexOrder.kernel_fcn_co_integral(t, α + 1, β)

resp = [y ye yi]

# Plots.plot(t, g)
Plots.plot(t, resp)


using AnalyticalEOR, Plots

swr = 0.2
sor = 0.2

kro0 = 1.0
krw0 = 0.5
nw = 3.0
no = 1.0

kr = RelPerms(swr, sor, krw0, kro0, nw, no)

μw = 1.0
μo = 3.0

si = 0.20
sj = 0.80

wf = solve_waterflooding(si, sj, kr, μw, μo)

tracer = solve_tracer(wf)

plot_tracer_fw(wf, tracer)

plot_sw_profile(wf, tracer, 0.1)
animate_sw_profile(wf, tracer)


# Mutation operators
# ==================

# Evolutionary mutations
# ======================

"""
    gaussian(x, s::IsotropicStrategy)

Performs Gaussian isotropic mutation of the recombinant `x` given the strategy `s`  by adding Gaussian noise as follows:

- ``x_i^\\prime = x_i + s.\\sigma \\mathcal{N}_i(0,1)``

"""
function gaussian(recombinant::AbstractVector, s::IsotropicStrategy)
    vals = randn(length(recombinant)) * s.σ
    recombinant += vals
    return recombinant
end

"""
    gaussian(x, s::AnisotropicStrategy)

Performs Gaussian anisotropic mutation of the recombinant `x` given the strategy `s`  by adding Gaussian noise as follows:

- ``x_i^\\prime = x_i + s.\\sigma_i \\mathcal{N}_i(0,1)``

"""
function gaussian(recombinant::AbstractVector, s::AnisotropicStrategy)
    @assert length(s.σ) == length(recombinant) "Parameter `σ` must be defined for every dimension of objective parameter"
    vals = randn(length(recombinant)) .* s.σ
    recombinant += vals
    return recombinant
end

"""
    cauchy(x, s::IsotropicStrategy)

Performs isotropic mutation of the recombinant `x` given the strategy `s`  by adding a noise from the Cauchy distribution as follows:

- ``x_i^\\prime = x_i + s.\\sigma_i \\delta_i``

where ``\\delta`` is a Cauchy random variable with the scale parameter ``t = 1`` [^2].

"""
function cauchy(recombinant::AbstractVector, s::IsotropicStrategy)
    l = length(recombinant)
    vals = s.σ * randn(l)./(randn(l).+eps())
    recombinant += vals
    return recombinant
end

# Strategy mutation operators
# ===========================

"""
    gaussian(s::IsotropicStrategy)

Performs in-place mutation of the isotropic strategy `s` modifying its mutated strategy parameter ``\\sigma`` with Gaussian noise as follows:

- ``\\sigma^\\prime = \\sigma \\exp(\\tau_0 \\mathcal{N}(0,1))``

"""
function gaussian(s::IsotropicStrategy)
    s.σ *= exp(s.τ₀*randn())
    return s
end

"""
    gaussian(s::AnisotropicStrategy)

Performs in-place mutation of the anisotropic strategy `s` modifying its mutated strategy parameter ``\\sigma`` with Gaussian noise as follows:

- ``\\sigma_i^\\prime = \\sigma_i \\exp(\\tau_0 \\mathcal{N}(0,1) + \\tau_i \\mathcal{N}(0,1))``

"""
function gaussian(s::AnisotropicStrategy)
    s.σ .*= exp.(s.τ₀*randn())*exp.(s.τ*randn(length(s.σ)))
    return s
end

# Genetic mutations
# =================

# Binary mutations
# ---------------------

"""
    flip(recombinant)

Returns an in-place mutated binary `recombinant` with a bit flips at random positions.
"""
function flip(recombinant::T) where {T <: BitVector}
    s = length(recombinant)
    pos = rand(1:s)
    recombinant[pos] = !recombinant[pos]
    return recombinant
end

"""
    bitinversion(recombinant)

Returns an in-place mutated binary `recombinant` with its bits inverted.
"""
bitinversion(recombinant::T) where {T <: BitVector} = map!(!, recombinant, recombinant)

# Real-valued mutations
# ---------------------

"""
    uniform(r = 1.0)

Returns an in-place real valued mutation function that performs the uniform distributed mutation [^1].

The mutation operator randomly chooses a number ``z`` in from the uniform distribution on the interval ``[-r,r]``, the mutation range.
The mutated individual is given by

- ``x_i^\\prime = x_i + z_i``

"""
function uniform(r::Real = 1.0)
    function mutation(recombinant::T) where {T <: AbstractVector}
        d = length(recombinant)
        recombinant += 2r.*rand(d).-r
        return recombinant
    end
    return mutation
end

"""
    gaussian(σ = 1.0)

Returns an in-place real valued mutation function that performs the normal distributed mutation [^1].

The mutation operator randomly chooses a number ``z`` in from the normal distribution ``\\mathcal{N}(0,\\sigma)`` with standard deviation ``\\sigma``.
The mutated individual is given by

- ``x_i^\\prime = x_i + z_i``

"""
function gaussian(σ::Real = 1.0)
    function mutation(recombinant::T) where {T <: AbstractVector}
        d = length(recombinant)
        recombinant += σ.*randn(d)
        return recombinant
    end
    return mutation
end

"""
    domainrange(valrange, m = 20)

Returns an in-place real valued mutation function that performs the BGA mutation scheme with the mutation range `valrange` and the mutation probability `1/m` [^1].
"""
function domainrange(valrange::Vector, m::Int = 20)
    prob = 1.0 / m
    function mutation(recombinant::T) where {T <: AbstractVector}
        d = length(recombinant)
        @assert length(valrange) == d "Range matrix must have $(d) columns"
        δ = zeros(m)
        for i in 1:length(recombinant)
            for j in 1:m
                δ[j] = (rand() < prob) ? δ[j] = 2.0^(-j) : 0.0
            end
            if rand() > 0.5
                recombinant[i] += sum(δ)*valrange[i]
            else
                recombinant[i] -= sum(δ)*valrange[i]
            end
        end
        return recombinant
    end
    return mutation
end

"""
    PM(lower, upper, p = 2)

Returns an in-place real valued mutation function that performs the Power Mutation (PM) scheme within `lower` and `upper` bound, and an index of
mutation `p`[^3].

*Note:* The implementation is a degenerate case of Mixed Integer Power Mutation ([`MIPM`](@ref))
"""
function PM(lower::Vector, upper::Vector, p::Float64 = 5.0) # index of distribution p
    return mipmmutation(lower, upper, p)
end

"""
    MIPM(lower, upper, p_real = 10, p_int = 4)

Returns an in-place real valued mutation function that performs the Mixed Integer Power Mutation (MI-PM) scheme within `lower` and `upper` bound, and an index of mutation `p_real` for real value and `p_int` for integer values[^4].
"""
function MIPM(lowerBounds::Vector, upperBounds::Vector, p_real::Float64 = 10.0, p_int::Float64 = 4.0) # index of distribution p
    return mipmmutation(lowerBounds, upperBounds, p_real, p_int)
end

function mipmmutation(lowerBounds::Vector, upperBounds::Vector, p_real::Float64, p_int::Union{Nothing, Float64} = nothing)
    function pm_mutation(x, l, u, s, d)
        x̄ = d < rand() ? x - s * (x - l) : x + s * (u - x)
        if isa(x, Integer)
            if isinteger(x̄)
                Int(x̄)
            else
                floor(Int, x̄) + (rand() > 0.5)
            end
        else
            x̄
        end
    end
    function mutation(recombinant::T) where {T <: Vector}
        d = length(recombinant)
        @assert length(lowerBounds) == d "Bounds vector must have $(d) columns"
        @assert length(upperBounds) == d "Bounds vector must have $(d) columns"
        @assert !(p_int === nothing && any(isa(x, Integer) for x in recombinant)) "Need to set p_int for integer variables"
        u = rand()
        P = (isa(x, Integer) ? p_int : p_real for x in recombinant)
        S = u .^ (1 ./ P) # random var following power distribution
        D = (recombinant - lowerBounds) ./ (upperBounds - lowerBounds)
        broadcast!(pm_mutation, recombinant, recombinant, lowerBounds, upperBounds, S, D)
        return recombinant
    end
    return mutation
end

# Combinatorial mutations (applicable to binary vectors)
# ------------------------------------------------------

"""
    inversion(recombinant)

Returns an in-place mutated individual with a random arbitrary length segment of the genome in the reverse order.
"""
function inversion(recombinant::T) where {T <: AbstractVector}
    l = length(recombinant)
    from, to = rand(1:l, 2)
    from, to = from > to ? (to, from)  : (from, to)
    l = round(Int,(to - from)/2)
    for i in 0:(l-1)
        swap!(recombinant, from+i, to-i)
    end
    return recombinant
end

function insertion(recombinant::T) where {T <: AbstractVector}
    l = length(recombinant)
    from, to = rand(1:l, 2)
    val = recombinant[from]
    deleteat!(recombinant, from)
    return insert!(recombinant, to, val)
end

"""
    swap2(recombinant)

Returns an in-place mutated individual with a two random elements of the genome are swapped.
"""
function swap2(recombinant::T) where {T <: AbstractVector}
    l = length(recombinant)
    from, to = rand(1:l, 2)
    swap!(recombinant, from, to)
    return recombinant
end

"""
    scramble(recombinant)

Returns an in-place mutated individual with elements, on a random arbitrary length segment of the genome, been scrambled.
"""
function scramble(recombinant::T) where {T <: AbstractVector}
    l = length(recombinant)
    from, to = rand(1:l, 2)
    from, to = from > to ? (to, from)  : (from, to)
    diff = to - from + 1
    if diff > 1
        patch = recombinant[from:to]
        idx = randperm(diff)
        # println("$(from)-$(to), P: $(patch), I: $(idx)")
        for i in 1:(diff)
            recombinant[from+i-1] = patch[idx[i]]
        end
    end
    return recombinant
end

"""
    shifting(recombinant)

Returns an in-place mutated individual with a random arbitrary length segment of the genome been shifted to an arbitrary position.
"""
function shifting(recombinant::T) where {T <: AbstractVector}
    l = length(recombinant)
    from, to, where = sort(rand(1:l, 3))
    patch = recombinant[from:to]
    diff = where - to
    if diff > 0
        # move values after tail of patch to the patch head position
        #println([from, to, where, diff])
        for i in 1:diff
            recombinant[from+i-1] = recombinant[to+i]
        end
        # place patch values in order
        start = from + diff
        for i in 1:length(patch)
            recombinant[start+i-1] = patch[i]
        end
    end
    return recombinant
end

# Differential Evolution
# ======================

"""
    differentiation(recombinant, mutators; F = 1.0)

Returns an in-place differentially mutated individual ``x^\\prime`` from `recombinant` ``x``  by `mutators` ``\\{\\xi_1, \\ldots, \\xi_n \\}`` as follows

- ``x^\\prime = x + \\sum_{i=1}^{n/2} F (\\xi_{2i-1} - \\xi_{2i})``

"""
function differentiation(recombinant::T, mutators::AbstractVector{T}; F::Real = 1.0) where {T <: AbstractVector}
    m = length(mutators)
    @assert m%2 == 0 "Must be even number of target mutators"
    for i in 1:2:m
        recombinant .+= F.*(mutators[i] .- mutators[i+1])
    end
    return recombinant
end

# Utils
# =====
function swap!(v::T, from::Int, to::Int) where {T <: AbstractVector}
    val = v[from]
    v[from] = v[to]
    v[to] = val
end

function mutationwrapper(gamutation::Function)
    wrapper(recombinant::T, s::S) where {T <: AbstractVector, S <: AbstractStrategy} =  gamutation(recombinant)
    return wrapper
end


# We have n particles here.
# For each particle i (i ∈ N⁺)  :
# We have Mₘᵢₙ ≤ mᵢ ≤ Mₘₐₓ
# We have an three dimentional vector v⃗ᵢ for i's velocity.
# And an three dimentional point Pᵢ for i's locationn.
# We have Gravitational Constant G ≈ 6.67408*10⁻¹¹ (m³kg⁻¹s⁻²) .
# And a small time step Δt→0 (actually, it's a very small value).
# At the beginning, every particle will spread in an area randomly.
# At the beginning, every particle will have a velocity which direction and size are all given randomly but in a range.
#
# ⁽ᶜᵒᵈᵉ ᵇʸ ᴿᵉˢᵇⁱ⁾

# Modules
using LinearAlgebra

# Constants
const n = 1024*16
const G = 6.67408e-11
const Δt = 60*12
const Δt² = Δt^2
const Δtₕ = Δt/2
const c = 299792458 # m/s
const Δt⁻²c⁻² = (c*Δt)^(-2)
const Mₘₐₓ = 1e22
const Mₘᵢₙ = 1e20
const Areaₘᵢₙ = [-1e11 -1e11 -1e11]
const Areaₘₐₓ = [1e11 1e11 1e11]
const Areaₜₑₘₚ = Areaₘₐₓ - Areaₘᵢₙ
const v⃗ₘᵢₙ  = [-5e2 -5e2 -5e2]
const v⃗ₘₐₓ  = [5e2 5e2 5e2]
const v⃗ₜₑₘₚ  = v⃗ₘₐₓ  - v⃗ₘᵢₙ
const SavePer = 5
const Total = 1024*16

# Generate
MGΔtΔtₕ = (rand(n)*(Mₘₐₓ-Mₘᵢₙ) + ones(n)*Mₘᵢₙ)*G*Δt*Δtₕ
P₀ = zeros(3, n)
P₁ = zeros(3, n)
Δtv⃗ₕ  = zeros(3, n)
#Δt²a⃗ₕ = zeros(3, n)
for i=1:3
  P₀[i, :] = rand(n)*Areaₜₑₘₚ[i] + ones(n)*Areaₘᵢₙ[i]
  Δtv⃗ₕ[i,  :] = (rand(n)*v⃗ₜₑₘₚ[i]  + ones(n)*v⃗ₘᵢₙ[i])*Δtₕ
end

#Set a blackhole
P₀[:, 1] = [0.0 0.0 0.0]
Δtv⃗ₕ[:, 1] = [0.0 0.0 0.0]
const MGΔtΔtₕ[1] = 1e30*G*Δt*Δtₕ

# Functions
function update(SYNC_ARRAY, P₀, P₁, Δtv⃗ₕ)
  @inbounds Threads.@threads for i=1:n
  #for i=1:n
    Δt²a⃗ₕ  = [0.0, 0.0, 0.0]
    #γ = sqrt(1 - (norm(Δtv⃗ₕ[i])^2)*Δt⁻²c⁻²)
    for j=1:(i-1)
      Δt²a⃗ₕ  += MGΔtΔtₕ[j]/(norm(P₀[:, j] - P₀[:, i])^3)*(P₀[:, j] - P₀[:, i])
    end
    for j=(i+1):n
      Δt²a⃗ₕ  += MGΔtΔtₕ[j]/(norm(P₀[:, j] - P₀[:, i])^3)*(P₀[:, j] - P₀[:, i])
    end
    P₁[:, i] = P₀[:, i] + 2*Δtv⃗ₕ[:, i] + Δt²a⃗ₕ #*γ
    Δtv⃗ₕ[:, i] = Δtv⃗ₕ[:, i] + Δt²a⃗ₕ #*γ
    # Mark
    #SYNC_ARRAY[i] = 1
  end
end

# Run
for j=0:Total-1
  SYNC_ARRAY = zeros(n)
  if mod(j, 2)==0
    @time update(SYNC_ARRAY, P₀, P₁, Δtv⃗ₕ)
  else
    @time update(SYNC_ARRAY, P₁, P₀, Δtv⃗ₕ)
  end
  # Wait
  #while sum(SYNC_ARRAY) < n
  #  pass
  #end
  if mod(j, SavePer)==0
    println(Δtv⃗ₕ[:, 1]/Δtₕ)
    println(P₀[:, 1])
  end
end


# finiteEnvironment.jl
# The environment of finite MPS and MPO form of Hamiltonian.
# an be used to do DMRG or calculating physical observables
"""
    struct FiniteEnvironment{S<:EuclideanSpace}
"""
struct FiniteEnvironment{S<:EuclideanSpace}
    left_env::Vector{<:AbstractTensorMap{S,1,2}}
    right_env::Vector{<:AbstractTensorMap{S,2,1}}
end

function initialize_FiniteEnvironment(mps::FiniteMPS, mpo::MPO, final_position::Int)
    spacetype(mps) == spacetype(mpo) || error("mps and mpo must have the same spacetype!")
    mps.canonical_form[1] == :right_canonical ||
        throw(ArgumentError("input mps should be in the right canonical form!"))
    N = length(mps.Ts)
    N == length(mpo.Ws) || error("mps and mpo must have the same length!")
    left_env = append!([TensorMap(undef, codomain(mps.Ts[i],1),
                            codomain(mps.Ts[i],1)*codomain(mpo.Ws[i],1)) for i in 1:N],
                       [TensorMap(undef, domain(mps.Ts[end]),
                            domain(mps.Ts[end])*domain(mpo.Ws[end],2))])
    right_env = append!([TensorMap(undef,codomain(mps.Ts[i],1)*codomain(mpo.Ws[i],1),
                            codomain(mps.Ts[i],1)) for i in 1:N],
                        [TensorMap(undef,domain(mps.Ts[end])*domain(mpo.Ws[end],2),
                            domain(mps.Ts[end]))])
    chi_mpo_left_boundary = dim(codomain(mpo.Ws[1],1))
    chi_mpo_right_boundary = dim(domain(mpo.Ws[end],2))
    left_boundary_matrix = zeros(1,1,chi_mpo_left_boundary)
    right_boundary_matrix = zeros(1,chi_mpo_right_boundary,1)
    left_boundary_matrix[1,1,final_position] = 1
    right_boundary_matrix[1,1,1] = 1
    left_boundary_tensor = TensorMap(left_boundary_matrix, codomain(mps.Ts[1],1),
                                        codomain(mps.Ts[1],1)*codomain(mpo.Ws[1],1))
    right_boundary_tensor = TensorMap(right_boundary_matrix,
                            domain(mps.Ts[end])*domain(mpo.Ws[end],2),domain(mps.Ts[end]))
    left_env[1] = left_boundary_tensor
    right_env[end] = right_boundary_tensor
    env = FiniteEnvironment{spacetype(mps)}(left_env,right_env)
    for i in N:-1:1
        update_right_env(i, mps, mpo, env)
    end
    return env
end

function update_left_env(i::Int, mps::FiniteMPS, mpo::MPO, env::FiniteEnvironment)
    @tensor A_tmp1[c d;b e] := env.left_env[i-1][c;a b]*mps.Ts[i-1][a d; e]
    @tensor A_tmp2[a e;d f] := A_tmp1[a b;c d]*mpo.Ws[i-1][c e;b f]
    @tensor A_tmp3[e;c d] := conj(mps.Ts[i-1][a b;e])*A_tmp2[a b;c d]
    env.left_env[i] = A_tmp3
end

function update_right_env(i::Int, mps::FiniteMPS, mpo::MPO, env::FiniteEnvironment)
    @tensor A_tmp1[d e b;c] := mps.Ts[i][d e;a]*env.right_env[i+1][a b;c]
    @tensor A_tmp2[a e f;d] := mpo.Ws[i][e f;b c]*A_tmp1[a b c;d]
    @tensor A_tmp3[a b;e] := A_tmp2[a b c;d]*conj(mps.Ts[i][e c;d])
    env.right_env[i] = A_tmp3
end


using JuMP
using Gurobi
using Suppressor
using CSV
using JLD
using Printf
using Clustering
gurobi_env = @suppress Gurobi.Env()


export RidgeRegressor, PCA

import MLJBase
import MLJBase: @mlj_model, metadata_model, metadata_pkg
# using Distances
using LinearAlgebra
using Tables
using ScientificTypes

import MultivariateStats

const MS = MultivariateStats

struct LinearFitresult{F} <: MLJBase.MLJType
    coefficients::Vector{F}
    intercept::F
end

const RIDGE_DESCR = "Ridge regressor with regularization parameter lambda. Learns a linear regression with a penalty on the l2 norm of the coefficients."
const PCA_DESCR = "Principal component analysis. Learns a linear transformation to project the data  on a lower dimensional space while preserving most of the initial variance."

####
#### RIDGE
####

"""
RidgeRegressor(; lambda=1.0)

$RIDGE_DESCR

## Parameters

* `lambda=1.0`: non-negative parameter for the regularization strength.
"""
@mlj_model mutable struct RidgeRegressor <: MLJBase.Deterministic
    lambda::Real = 1.0::(_ ≥ 0)
end

function MLJBase.fit(model::RidgeRegressor, verbosity::Int, X, y)
    Xmatrix   = MLJBase.matrix(X)
    features  = Tables.schema(X).names
    θ         = MS.ridge(Xmatrix, y, model.lambda)
    coefs     = θ[1:end-1]
    intercept = θ[end]

    fitresult = LinearFitresult(coefs, intercept)
    report    = NamedTuple()
    cache     = nothing

    return fitresult, cache, report
end

MLJBase.fitted_params(::RidgeRegressor, fr) =
    (coefficients=fr.coefficients, intercept=fr.intercept)

function MLJBase.predict(::RidgeRegressor, fr, Xnew)
    Xmatrix = MLJBase.matrix(Xnew)
    return Xmatrix * fr.coefficients .+ fr.intercept
end

####
#### PCA
####

const PCAFitResultType = MS.PCA

"""
PCA(; maxoutdim=nothing, method=:auto, pratio=0.99, mean=nothing)

$PCA_DESCR

## Parameters

* `maxoutdim=nothing`: maximum number of output dimensions, unconstrained if nothing.
* `method=:auto`: method to use to solve the problem, one of `:auto`, `:cov` or `:svd`
* `pratio=0.99`: ratio of variance preserved
* `mean=nothing`: if set to nothing centering will be computed and applied, if set to `0` no centering (assumed pre-centered), if a vector is passed, the centering is done with that vector.
"""
@mlj_model mutable struct PCA <: MLJBase.Unsupervised
    maxoutdim::Union{Nothing,Int} = nothing::(_ === nothing || _ ≥ 1)
    method::Symbol  = :auto::(_ in (:auto, :cov, :svd))
    pratio::Float64 = 0.99::(0.0 < _ ≤ 1.0)
    mean::Union{Nothing, Real, Vector{Float64}} = nothing::(_ === nothing || (_ isa Real && iszero(_)) || true)
end

function MLJBase.fit(model::PCA, verbosity::Int, X)
    Xarray = MLJBase.matrix(X)
    mindim = minimum(size(Xarray))

    maxoutdim = model.maxoutdim === nothing ? mindim : model.maxoutdim

    # NOTE: copy/transpose
    fitresult = MS.fit(MS.PCA, permutedims(Xarray);
                       method=model.method,
                       pratio=model.pratio,
                       maxoutdim=maxoutdim,
                       mean=model.mean)

    cache = nothing
    report = (indim=size(fitresult)[1],
              outdim=size(fitresult)[2],
              mean=MS.mean(fitresult),
              principalvars=MS.principalvars(fitresult),
              tprincipalvar=MS.tprincipalvar(fitresult),
              tresidualvar=MS.tresidualvar(fitresult),
              tvar=MS.var(fitresult))

    return fitresult, cache, report
end

MLJBase.fitted_params(::PCA, fr) = (projection=fr,)

function MLJBase.transform(::PCA, fr::PCAFitResultType, X)
    # X is n x d, need to transpose and copy twice...
    Xarray = MLJBase.matrix(X)
    Xnew   = permutedims(MS.predict(fr, permutedims(Xarray)))
    return MLJBase.table(Xnew, prototype=X)
end

####
#### METADATA
####

metadata_pkg.((RidgeRegressor, PCA),
              name="MultivariateStats",
              uuid="6f286f6a-111f-5878-ab1e-185364afe411",
              url="https://github.com/JuliaStats/MultivariateStats.jl",
              license="MIT",
              julia=true,
              is_wrapper=false)

metadata_model(RidgeRegressor,
               input=MLJBase.Table(MLJBase.Continuous),
               target=AbstractVector{MLJBase.Continuous},
               weights=false,
               descr=RIDGE_DESCR)

metadata_model(PCA,
               input=MLJBase.Table(MLJBase.Continuous),
               target=MLJBase.Table(MLJBase.Continuous),
               weights=false,
               descr=PCA_DESCR)


# Forming adjacency graphs.

# Read CSV file into associative array.
E = ReadCSV(joinpath(Base.source_dir(),"Edge.csv"))

# Get vertices and convert to numbers.
Ev = logical( E[:, StartsWith("V,")] )

# Compute vertex adjacency graph.
Av = sqIn(Ev)
printFull(Av)

# Compute edge adjacency graph.
Ae = sqOut(Ev)
printFull(Ae)


"""Potential vorticity calculated as q = (f + ∂v/∂x - ∂u/∂y)/h."""
function PV!(Diag::DiagnosticVars,S::ModelSetup)

    @unpack q,dvdx,dudy,h_q = Diag.Vorticity
    @unpack f_q,ep = S.grid

    m,n = size(q)
    @boundscheck (m,n) == size(f_q) || throw(BoundsError())
    @boundscheck (m+2,n+2) == size(dvdx) || throw(BoundsError())
    @boundscheck (m+2+ep,n+2) == size(dudy) || throw(BoundsError())
    @boundscheck (m,n) == size(h_q) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            q[i,j] = (f_q[i,j] + dvdx[i+1,j+1] - dudy[i+1+ep,j+1]) / h_q[i,j]
        end
    end
end

"""Transit function to call the specified advection scheme."""
function PVadvection!(Diag::DiagnosticVars,S::ModelSetup)
    if S.parameters.adv_scheme == "Sadourny"
        PV_Sadourny!(Diag)
    elseif S.parameters.adv_scheme == "ArakawaHsu"
        PV_ArakawaHsu!(Diag)
    end
end

"""Advection of potential vorticity qhv,qhu as in Sadourny, 1975
enstrophy conserving scheme."""
function PV_Sadourny!(Diag::DiagnosticVars)

    @unpack U,V = Diag.VolumeFluxes
    @unpack V_u,U_v = Diag.Vorticity
    @unpack q_u,q_v,qhv,qhu,ep = Diag.Vorticity

    Ixy!(V_u,V)
    Ixy!(U_v,U)

    m,n = size(qhv)
    @boundscheck (m+2-ep,n) == size(q_u) || throw(BoundsError())
    @boundscheck (m+2-ep,n) == size(V_u) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            qhv[i,j] = q_u[i+1-ep,j] * V_u[i+1-ep,j]
        end
    end

    m,n = size(qhu)
    @boundscheck (m,n+2) == size(q_v) || throw(BoundsError())
    @boundscheck (m,n+2) == size(U_v) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            qhu[i,j] = q_v[i,j+1] * U_v[i,j+1]
        end
    end
end

"""Advection of potential vorticity qhv,qhu as in Arakawa and Hsu, 1990
Energy and enstrophy conserving (in the limit of non-divergent mass flux) scheme with τ = 0."""
function PV_ArakawaHsu!(Diag::DiagnosticVars)

    @unpack U,V = Diag.VolumeFluxes
    @unpack qhv,qhu,ep = Diag.Vorticity
    @unpack qα,qβ,qγ,qδ = Diag.ArakawaHsu

    # Linear combinations of q and V=hv to yield qhv
    m,n = size(qhv)
    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            qhv[i,j] = qα[1-ep+i,j]*V[2-ep+i,j+1] + qβ[1-ep+i,j]*V[1-ep+i,j+1] + qγ[1-ep+i,j]*V[1-ep+i,j] + qδ[1-ep+i,j]*V[2-ep+i,j]
        end
    end

    # Linear combinations of q and U=hu to yield qhu
    m,n = size(qhu)
    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            qhu[i,j] = qα[i,j]*U[i,j+1] + qβ[i+1,j]*U[i+1,j+1] + qγ[i+1,j+1]*U[i+1,j+2] + qδ[i,j+1]*U[i,j+2]
        end
    end
end

""" Linear combination α of potential voriticity q according
to the energy and enstrophy conserving scheme of Arakawa and Hsu, 1990"""
function AHα!(α::Array{T,2},q::Array{T,2}) where {T<:AbstractFloat}
    m,n = size(α)
    @boundscheck (m+1,n+1) == size(q) || throw(BoundsError())

    one_twelve = T(1/12)

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            α[i,j] = one_twelve*(q[i,j] + q[i,j+1] + q[i+1,j+1])
        end
    end
end

""" Linear combination δ of potential voriticity q according
to the energy and enstrophy conserving scheme of Arakawa and Hsu, 1990 """
function AHβ!(β::Array{T,2},q::Array{T,2}) where {T<:AbstractFloat}
    m,n = size(β)
    @boundscheck (m,n+1) == size(q) || throw(BoundsError())

    one_twelve = T(1/12)

    @inbounds for j ∈ 1:n
        β[1,j] = one_twelve*(q[1,j] + q[1,j+1] + q[end,j+1])
        for i ∈ 2:m
            β[i,j] = one_twelve*(q[i,j] + q[i,j+1] + q[i-1,j+1])
        end
    end
end

""" Linear combination γ of potential voriticity q according
to the energy and enstrophy conserving scheme of Arakawa and Hsu, 1990 """
function AHγ!(γ::Array{T,2},q::Array{T,2}) where {T<:AbstractFloat}
    m,n = size(γ)
    @boundscheck (m,n+1) == size(q) || throw(BoundsError())

    one_twelve = T(1/12)

    @inbounds for j ∈ 1:n
        γ[1,j] = one_twelve*(q[1,j] + q[1,j+1] + q[end,j])
        for i ∈ 2:m
            γ[i,j] = one_twelve*(q[i,j] + q[i-1,j] + q[i,j+1])
        end
    end
end

""" Linear combination β of potential voriticity q according
to the energy and enstrophy conserving scheme of Arakawa and Hsu, 1990 """
function AHδ!(δ::Array{T,2},q::Array{T,2}) where {T<:AbstractFloat}
    m,n = size(δ)
    @boundscheck (m+1,n+1) == size(q) || throw(BoundsError())

    one_twelve = T(1/12)

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            δ[i,j] = one_twelve*(q[i,j] + q[i,j+1] + q[i+1,j])
        end
    end
end


import ..UncertainValues.AbstractUncertainValue
import ..SamplingConstraints:
    SamplingConstraint,
	NoConstraint,
	TruncateLowerQuantile,
	TruncateUpperQuantile,
	TruncateQuantiles,
	TruncateMinimum,
	TruncateMaximum,
	TruncateRange,
    TruncateStd
import Distributions: truncated

########################################################################
# Resampling with constraints
########################################################################

"""
Resample an uncertain value by first truncating its furnishing distribution with the
provided sampling `constraint`.
"""
function resample(uv::AbstractUncertainValue, constraint::SamplingConstraint) end

"""
    resample(uv::AbstractUncertainValue, constraint::NoConstraint)

Resample an uncertain value without contraints (use the full furnishing distribution).

## Example

```julia
uncertainval = UncertainValue(0, 0.2, Normal)

# Resample the uncertain value by resampling the full distribution once.
resample(uncertainval, NoConstraint())
```
"""
resample(uv::AbstractUncertainValue, constraint::NoConstraint) = resample(uv)

"""
    resample(uv::AbstractUncertainValue, constraint::NoConstraint, n::Int)

Resample an uncertain value without contraints (use the full furnishing distribution).

## Example

```julia
uncertainval = UncertainValue(0, 0.2, Normal)

# Resample the uncertain value by resampling the full distribution 1000 times.
resample(uncertainval, NoConstraint(), 1000)
```
"""
resample(uv::AbstractUncertainValue, constraint::NoConstraint, n::Int) = resample(uv, n)

"""
    resample(uv::AbstractUncertainValue, constraint::TruncateUpperQuantile)

Resample by first truncating the distribution representing the value at a lower
quantile, then performing the resampling.

## Example

```julia
uncertainval = UncertainValue(0, 0.2, Normal)
constraint = TruncateLowerQuantile(0.16)

# Resample the uncertain value by truncating the distribution furnishing it,
# then resampling the new distribution once.
resample(uncertainval, constraint)
```
"""
function resample(uv::AbstractUncertainValue, constraint::TruncateLowerQuantile)
    # Apply (another level of) truncation, then sample
    lower_bound = quantile(uv.distribution, constraint.lower_quantile)
    upper_bound = maximum(uv)
    rand(truncated(uv.distribution, lower_bound, upper_bound))
end

"""
    resample(uv::AbstractUncertainValue, constraint::TruncateLowerQuantile, n::Int)

Resample by first truncating the distribution representing the value at a lower
quantile, then performing the resampling.

## Example

```julia
uncertainval = UncertainValue(0, 0.2, Normal)
constraint = TruncateLowerQuantile(0.16)

# Resample the uncertain value by truncating the distribution furnishing it,
# then resampling the new distribution 1000 times.
resample(uncertainval, constraint, 1000)
```
"""
function resample(uv::AbstractUncertainValue, constraint::TruncateLowerQuantile, n::Int)
    # Apply (another level of) truncation, then sample
    lower_bound = quantile(uv.distribution, constraint.lower_quantile)
    upper_bound = maximum(uv)
    rand(truncated(uv.distribution, lower_bound, upper_bound), n)
end

"""
    resample(uv::AbstractUncertainValue, constraint::TruncateUpperQuantile)

Resample by first truncating the distribution representing the value at an upper
quantile, then performing the resampling.

## Example

```julia
uncertainval = UncertainValue(0, 0.2, Normal)
constraint = TruncateUpperQuantile(0.8)

# Resample the uncertain value by truncating the distribution furnishing it,
# then resampling the new distribution once.
resample(uncertainval, constraint)
```
"""
function resample(uv::AbstractUncertainValue, constraint::TruncateUpperQuantile)
    # Apply (another level of) truncation, then sample
    lower_bound = minimum(uv)
    upper_bound = quantile(uv.distribution, constraint.upper_quantile)
    rand(truncated(uv.distribution, lower_bound, upper_bound))
end

"""
    resample(uv::AbstractUncertainValue, constraint::TruncateUpperQuantile, n::Int)

Resample by first truncating the distribution representing the value at an upper
quantile, then performing the resampling.

## Example

```julia
uncertainval = UncertainValue(0, 0.2, Normal)
constraint = TruncateUpperQuantile(0.8)

# Resample the uncertain value by truncating the distribution furnishing it,
# then resampling the new distribution 1000 times.
resample(uncertainval, constraint, 1000)
```
"""
function resample(uv::AbstractUncertainValue, constraint::TruncateUpperQuantile, n::Int)
    # Apply (another level of) truncation, then sample
    lower_bound = minimum(uv)
    upper_bound = quantile(uv.distribution, constraint.upper_quantile)
    rand(truncated(uv.distribution, lower_bound, upper_bound), n)
end

"""
    resample(uv::AbstractUncertainValue, constraint::TruncateQuantiles)

Resample by first truncating the distribution representing the value at a set of
qunatiles, then performing the resampling.

## Example

```julia
uncertainval = UncertainValue(0, 1, Uniform)
constraint = TruncateLowerQuantile(0.2)

# Resample the uncertain value by truncating the distribution furnishing it,
# then resampling the new distribution once.
resample(uncertainval, constraint)
```
"""
function resample(uv::AbstractUncertainValue, constraint::TruncateQuantiles)
    # Apply (another level of) truncation, then sample
    lower_bound = quantile(uv.distribution, constraint.lower_quantile)
    upper_bound = quantile(uv.distribution, constraint.upper_quantile)
    rand(truncated(uv.distribution, lower_bound, upper_bound))
end

"""
    resample(uv::AbstractUncertainValue, constraint::TruncateQuantiles, n::Int)

Resample by first truncating the distribution representing the value at a set of
qunatiles, then performing the resampling.

## Example

```julia
uncertainval = UncertainValue(0, 1, Uniform)
constraint = TruncateLowerQuantile(0.2)

# Resample the uncertain value by truncating the distribution furnishing it,
# then resampling the new distribution 1000 times.
resample(uncertainval, constraint, 1000)
```
"""
function resample(uv::AbstractUncertainValue, constraint::TruncateQuantiles, n::Int)
    # Apply (another level of) truncation, then sample
    lower_bound = quantile(uv.distribution, constraint.lower_quantile)
    upper_bound = quantile(uv.distribution, constraint.upper_quantile)
    rand(truncated(uv.distribution, lower_bound, upper_bound), n)
end

"""
    resample(uv::AbstractUncertainValue, constraint::TruncateStd, n::Int;
        n_draws::Int = 10000)

Resample by first truncating the distribution representing the value to some multiple
of its standard deviation around the mean.

## Example

```julia
uncertainval = UncertainValue(0, 0.8, Normal)
constraint = TruncateStd(1.1) # accept values only in range 1.1*stdev around the mean

# Resample the uncertain value by truncating the distribution furnishing it,
# then resampling the new distribution 1000 times.
resample(uncertainval, constraint, 1000)
```
"""
function resample(uv::AbstractUncertainValue, constraint::TruncateStd, n::Int;
        n_draws::Int = 10000)
    # Apply (another level of) truncation, then sample
    stdev = std(resample(uv, n_draws))
    m = mean(resample(uv, n_draws))
    lower_bound = m - stdev
    upper_bound = m + stdev

    rand(truncated(uv.distribution, lower_bound, upper_bound), n)
end

"""
    resample(uv::AbstractUncertainValue, constraint::TruncateMinimum)

Resample by first truncating the distribution representing the value at some
minimum value, then performing the resampling.

## Example

```julia
uncertainval = UncertainValue(0, 0.2, Normal)
constraint = TruncateMinimum(-0.5) # accept no values less than -0.5

# Resample the uncertain value by truncating the distribution furnishing it,
# then resampling the new distribution once.
resample(uncertainval, constraint)
```
"""
function resample(uv::AbstractUncertainValue, constraint::TruncateMinimum)
    # Apply (another level of) truncation, then sample
    lower_bound = constraint.min
    upper_bound = maximum(uv)
    lower_bound > upper_bound ? error("lower bound > upper_bound") : nothing
    rand(truncated(uv.distribution, lower_bound, upper_bound))
end

"""
    resample(uv::AbstractUncertainValue, constraint::TruncateMinimum, n::Int)

Resample by first truncating the distribution representing the value at some
minimum value, then performing the resampling.

## Example

```julia
uncertainval = UncertainValue(0, 0.2, Normal)
constraint = TruncateMinimum(-0.5)

# Resample the uncertain value by truncating the distribution furnishing it,
# then resampling the new distribution 1000 times.
resample(uncertainval, constraint, 1000)
```
"""
function resample(uv::AbstractUncertainValue, constraint::TruncateMinimum, n::Int)
    # Apply (another level of) truncation, then sample
    lower_bound = constraint.min
    upper_bound = maximum(uv)
    lower_bound > upper_bound ? error("lower bound > upper_bound") : nothing
    rand(truncated(uv.distribution, lower_bound, upper_bound), n)
end

"""
    resample(uv::AbstractUncertainValue, constraint::TruncateMaximum)

Resample by first truncating the distribution representing the value at some
minimum value, then performing the resampling.

## Example

```julia
uncertainval = UncertainValue(0, 0.8, Normal)
constraint = TruncateMaximum(1.1) # accept no values larger than 1.1

# Resample the uncertain value by truncating the distribution furnishing it,
# then resampling the new distribution once.
resample(uncertainval, constraint)
```
"""
function resample(uv::AbstractUncertainValue, constraint::TruncateMaximum)
    # Apply (another level of) truncation, then sample
    upper_bound = constraint.max
    lower_bound = minimum(uv)
    lower_bound > upper_bound ? error("lower bound > upper_bound") : nothing
    rand(truncated(uv.distribution, lower_bound, upper_bound))
end

"""
    resample(uv::AbstractUncertainValue, constraint::TruncateMaximum, n::Int)

Resample by first truncating the distribution representing the value at some
minimum value, then performing the resampling.

## Example

```julia
uncertainval = UncertainValue(0, 0.8, Normal)
constraint = TruncateMaximum(1.1) # accept no values larger than 1.1

# Resample the uncertain value by truncating the distribution furnishing it,
# then resampling the new distribution 1000 times.
resample(uncertainval, constraint, 1000)
```
"""
function resample(uv::AbstractUncertainValue, constraint::TruncateMaximum, n::Int)
    # Apply (another level of) truncation, then sample
    lower_bound = minimum(uv)
    upper_bound = constraint.max
    lower_bound > upper_bound ? error("lower bound > upper_bound") : nothing
    rand(truncated(uv.distribution, lower_bound, upper_bound), n)
end

"""
    resample(uv::AbstractUncertainValue, constraint::TruncateRange)

Resample by first truncating the distribution representing the value at some
minimum value, then performing the resampling.

## Example

```julia
uncertainval = UncertainValue(0, 0.8, Normal)
constraint = TruncateRange(-0.7, 1.1) # accept values only in range [-0.7, 1.1]

# Resample the uncertain value by truncating the distribution furnishing it,
# then resampling the new distribution once.
resample(uncertainval, constraint)
```
"""
function resample(uv::AbstractUncertainValue, constraint::TruncateRange)
    # Apply (another level of) truncation, then sample
    lower_bound = constraint.min
    upper_bound = constraint.max
    lower_bound > upper_bound ? error("lower bound > upper_bound") : nothing

    rand(truncated(uv.distribution, lower_bound, upper_bound))
end

"""
    resample(uv::AbstractUncertainValue, constraint::TruncateRange, n::Int)

Resample by first truncating the distribution representing the value at some
minimum value, then performing the resampling.

## Example

```julia
uncertainval = UncertainValue(0, 0.8, Normal)
constraint = TruncateRange(-0.7, 1.1) # accept values only in range [-0.7, 1.1]

# Resample the uncertain value `1000` times by truncating the distribution furnishing it,
# then resampling the new distribution `1000` times
resample(uncertainval, constraint, 1000)
```
"""
function resample(uv::AbstractUncertainValue, constraint::TruncateRange, n::Int)
    # Apply (another level of) truncation, then sample
    lower_bound = constraint.min
    upper_bound = constraint.max
    lower_bound > upper_bound ? error("lower bound > upper_bound") : nothing

    rand(truncated(uv.distribution, lower_bound, upper_bound), n)
end

"""
    resample(uv::AbstractUncertainValue, constraint::TruncateStd; n_draws::Int = 10000)

Resample by first truncating the distribution representing the value to some multiple
of its standard deviation around the mean.

## Example

```julia
uncertainval = UncertainValue(0, 0.8, Normal)
constraint = TruncateStd(1.1) # accept values only in range 1.1*stdev around the mean

# Resample the uncertain value by truncating the distribution furnishing it,
# then resampling the new distribution 1000 times.
resample(uncertainval, constraint, 1000)
```
"""
function resample(uv::AbstractUncertainValue, constraint::TruncateStd; n_draws::Int = 10000)
    # Apply (another level of) truncation, then sample
    stdev = std(resample(uv, n_draws))
    m = mean(resample(uv, n_draws))
    lower_bound = m - stdev
    upper_bound = m + stdev
    lower_bound > upper_bound ? error("lower bound > upper_bound") : nothing
    rand(truncated(uv.distribution, lower_bound, upper_bound))
end

export resample


module MatrixFactorisation

using DataModel
using MFUtilities
using c_wrappers
using MLUtilities
using Distributions

export MatrixFactorisationModel

"""
This model implements a Bayesian matrix factorisation model for collaborative filtering following

Ahn et al. Large-Scale Distributed Bayesian Matrix Factorization using Stochastic Gradient MCMC

Ratings are modelled as Rᵢⱼ- mean_rating ∼ N(dot(Uᵢ,Vⱼ)+aᵢ+bⱼ,τ) where U,V,a,b are assumed to be normally distributed with mean 0 and precisions λᵤ (one per dim),λᵥ (one per dim),λₐ,λᵦ with Γ(α₀,β₀) hyperpriors.

Fields:
    - block: data in the format Nx3; first column user id between 1 and nusers; second column item id between 1 and nitems; third column rating.
    - nusers
    - nitems
    - d: dimension of the user and item feature vectors
    - τ: noise variance
    - α₀,β₀:  hyperprior parameters
    - n_elem_row: number of occurences per user
    - n_elem_col: number of occurences per item
    - inv_pr_pick_users, inv_pr_pick_items: needed for specialized sampler from Ahn et al.
    - batchiterator, batchstate: provides minibatches
    - ntrain: size of dataset
    - batchsize
    - mean_rating
"""
type MatrixFactorisationModel <: AbstractDataModel
    block::Array{Float64}
    nusers::Int64
    nitems::Int64
    d::Int64
    τ::Float64
    α₀::Float64
    β₀::Float64
    n_elem_row
    n_elem_col
    inv_pr_pick_users::Array{Float64}
    inv_pr_pick_items::Array{Float64}
    batchiterator::BatchIterator
    batchstate
    ntrain::Float64 # number of ratings in this block
    batchsize::Int64
    mean_rating::Float64
    function MatrixFactorisationModel(block,d,batchsize;τ= 2,α₀=1.,β₀=1., randomize=true)
        # get user and items numbers
        nusers    =  round(Int64,maximum(block[:,1]))
        nitems    = round(Int64,maximum(block[:,2]))

        # get data size
        ntrain = size(block)[1]

        # get number of user and item occurrence
        n_elem_row, n_elem_col = get_n_elem(block, nusers,nitems)

        # get auxiliary information about
        inv_pr_pick_users, inv_pr_pick_items = comp_inv_pr_pick(n_elem_row,n_elem_col,batchsize)

        # set up batch iterator
        batchiterator = BatchIterator(ntrain,batchsize,Inf,randomize=randomize)
        batchstate = start(batchiterator)

        mean_rating = mean(block[:,3])

        if randomize
            block = block[randperm(ntrain),:]
        end

        new(block,nusers,nitems,d,τ,α₀,β₀,n_elem_row,n_elem_col,inv_pr_pick_users,inv_pr_pick_items,batchiterator,batchstate,ntrain,batchsize,mean_rating)
    end
end

"""
fetchnparams returns the total number of parameters in the model.
"""
fetchnparams(dm::MatrixFactorisationModel) = (dm.nusers+dm.nitems)*(dm.d+1)

"""
getbatch returns the next minibatch
"""
function getbatch(dm::MatrixFactorisationModel)
    idx, dm.batchstate = next(dm.batchiterator,dm.batchstate)
    batch = dm.block[idx,:]
    return idx, batch
end

"""
getparams returns views of the user feature matrix U, the user biases a, the item feature matrix V and the item biases b
"""
function getparams(x::Array{Float64},dm::MatrixFactorisationModel)
    # D dx1
    # U n_u x d
    # a n_u x 1
    # V n_i x d
    # b n_i x 1

    d = dm.d
    nusers = dm.nusers
    nitems = dm.nitems

    ustart = 1
    uend = d*nusers
    astart = uend + 1
    aend = uend + nusers
    vstart = aend + 1
    vend = aend + d*nitems
    bstart = vend + 1
    bend = vend + nitems
    U = reshape(view(x,ustart:uend),(d,nusers))
    a = view(x,astart:aend)
    V = reshape(view(x,vstart:vend),(d,nitems))
    b = view(x,bstart:bend)
    return U,a,V,b
end

"""
getparams_λ returns views of the
"""
function getparams_λ(λ::Array{Float64},dm::MatrixFactorisationModel)
    d = dm.d

    @assert length(λ) == 2d+2
    λUstart = 1
    λUend = d
    λVstart = d+ 1
    λVend = 2d
    λaid = 2d+1
    λbid = 2d+2
    λU = view(λ,λUstart:λUend)
    λa = view(λ,λaid)
    λV = view(λ,λVstart:λVend)
    λb = view(λ,λbid)
    return λU,λV,λa,λb
end

"""
getgrad returns a function that gives the next stochastic gradient
"""
function DataModel.getgrad(dm::MatrixFactorisationModel;λ::Array{Float64,1}=ones(2*dm.d+2))
    function mf_gradient(x)
        err = Array(Float64,dm.batchsize);
        U,a,V,b = getparams(x,dm);
        idx, batch = getbatch(dm)
        rr = batch[:,3];
        uu = batch[:,1];
        ii = batch[:,2];

        ux, ix     = unique(uu), unique(ii)
        c_wrappers.c_comp_error!(err,rr,dm.mean_rating,U,V,a,b,uu,ii,dm.batchsize,dm.d);

        grad = zeros(fetchnparams(dm))
        gradU,grada,gradV,gradb = getparams(grad,dm)
        c_wrappers.c_comp_grad_sum!(err, uu, ii, U, V, a, b, gradU, gradV, grada, gradb, dm.batchsize, dm.d)

        return grad
    end
    return x -> dm.ntrain/dm.batchsize*mf_gradient(x) + gradprior(dm,x,λ)
end

"""
getllik returns a function that gives the next stochastic loglikehood estimate
"""
function DataModel.getllik(dm::MatrixFactorisationModel;λ::Array{Float64,1}=ones(2*dm.d+2))
    function loglikehood(x)
        err = Array(Float64,dm.batchsize);
        U,a,V,b = getparams(x,dm)
        λᵤ,λᵥ,λₐ,λᵦ = getparams_λ(λ,dm)
        idx, batch = getbatch(dm)
        rr = batch[:,3];
        uu = batch[:,1];
        ii = batch[:,2];

        c_wrappers.c_comp_error!(err,rr,dm.mean_rating,U,V,a,b,uu,ii,dm.batchsize,dm.d);
        logprior = -(fetchnparams(dm))/2*log(2*pi)+0.5*sum(log(λ))-0.5*(sum(λᵤ.*sum(U.*U,2))) - 0.5*(sum(λᵥ.*sum(V.*V,2))) - 0.5*(sum(λₐ*sum(a.*a))) - 0.5*(sum(λᵦ.*sum(b.*b)))
        llik = logprior - dm.ntrain/dm.batchsize*sum(0.5*err.^2) - dm.ntrain/2*log(2*π)

        return llik[1]
    end
    return loglikehood
end

"""
auxiliary function to calculate the prior gradient
"""
function gradprior(dm::MatrixFactorisationModel,x::Array{Float64}, λ::Array{Float64})

    λᵤ,λᵥ,λₐ,λᵦ = getparams_λ(λ,dm)

    U,a,V,b = getparams(x,dm)
    ∇ = zeros(fetchnparams(dm))
    ∇ᵤ,∇ₐ,∇ᵥ,∇ᵦ = getparams(∇,dm)
    ∇ᵤ[:] = (-λᵤ.*U)[:]
    ∇ᵥ[:] = (-λᵥ.*V)[:]
    ∇ₐ[:] = -λₐ.*(a)
    ∇ᵦ[:] = -λᵦ.*(b)
    return ∇
end

"""
evaluate returns the RMSE given parameters `x`
"""
function evaluate(dm::MatrixFactorisationModel,
  x::Vector{Float64}
  )
  prediction = predict(dm,x)
  evaluateprediction(dm,prediction)
end

"""
evaluate returns the RMSE given parameters `x` for a test model. Since some users/items might not be present in the test set, nusers/nitems has to be supplied to allow the parameters to be matched correctly.
"""
function evaluate_test(dm::MatrixFactorisationModel,
  x::Vector{Float64}, dmtrain::MatrixFactorisationModel)
  prediction = predict_test(dm,x,dmtrain)
  evaluateprediction(dm,prediction)
end

"""
evaluateprediction return the RMSE given predictions
"""
function evaluateprediction(dm::MatrixFactorisationModel,
  prediction::Vector{Float64})
  diff = dm.block[:,3] - prediction
  sse = sum(diff.*diff)
  rmse = sqrt(sse./dm.ntrain)
  Dict(:rmse => rmse)
end

"""
predict gives predicted ratings given parameters `x`
"""
function predict(dm::MatrixFactorisationModel,x::Vector{Float64})
    U,a,V,b = getparams(x,dm)
    prediction = zeros(round(Int64,dm.ntrain))
    c_wrappers.c_predict!(dm.block,U,V,a,b,prediction,1.,5.,dm.mean_rating)
    prediction
end

"""
predict gives predicted ratings given parameters `x`
"""
function predict_test(dm::MatrixFactorisationModel,x::Vector{Float64},dmtrain::MatrixFactorisationModel)
    U,a,V,b = getparams(x,dmtrain)
    prediction = zeros(round(Int64,dm.ntrain))
    c_wrappers.c_predict!(dm.block,U,V,a,b,prediction,1.,5.,dm.mean_rating)
    prediction
end

"""
lambda_sample performs a Gibbs update for the precisions λ
"""
function lambda_sample(dm::MatrixFactorisationModel,x::Array{Float64,1})
    U,a,V,b = getparams(x,dm)
    d = dm.d
    nusers  = dm.nusers
    nitems = dm.nitems
    α₀ = dm.α₀
    β₀ = dm.β₀

    αᵤ = α₀ + dm.nusers/2
    αₐ = αᵤ
    αᵥ = α₀ + dm.nitems/2
    αᵦ  = αᵥ

    βᵤ = β₀ + 0.5.*sum(U.*U,2)[:]
    βᵥ = β₀ + 0.5.*sum(V.*V,2)[:]
    βₐ = β₀ + 0.5.*sum(a.*a)[1]
    βᵦ = β₀ + 0.5.*sum(b.*b)[1]

    λᵤ = sample_gamma(ones(d)*αᵤ,βᵤ)
    λᵥ = sample_gamma(ones(d)*αᵥ,βᵥ)
    λₐ = sample_gamma(αₐ,βₐ)
    λᵦ = sample_gamma(αᵦ,βᵦ)
    return [λᵤ;λₐ;λᵥ;λᵦ]
end

"""
auxiliary function to sample gamma variables
"""
function sample_gamma(α::Real,β::Real)
    try
        @assert α > 0.0 && β > 0.0
    catch
        @show α,β
    end
    d = Gamma(α ,1/β)
    return Base.rand(d)
end

"""
auxiliary function to sample gamma variables
"""
function sample_gamma(α::Array{Float64},β::Array{Float64})
    @assert size(α) == size(β)
    return Float64[sample_gamma(a,b) for (a,b) in zip(α,β)]
end

"""
Specialised sparse sampler from Ahn et al.
"""
type SparseSGLDState
    x::Array{Float64,1}
    grad_sum::Array{Float64,1}
    λ::Array{Float64,1}
    t::Int
    stepsize::Function
    niters::Int
    function SparseSGLDState(x::Array{Float64,1},lambda::Array{Float64,1},stepsize::Function; niters::Int = 10)
        grad_sum = zeros(length(x))
        new(x,grad_sum,lambda,0,stepsize,niters)
    end
    function SparseSGLDState(x::Array{Float64,1},lambda::Array{Float64,1},stepsize::Float64; niters::Int= 10)
        grad_sum = zeros(length(x))
        f = niters::Int -> stepsize
        new(x,grad_sum,lambda,0,f,niters)
    end
end

"""
one update using the sparse SGLD sampler.
"""
function sample_sparse!(dm::MatrixFactorisationModel,
   state::SparseSGLDState)
        factor = dm.ntrain/(dm.batchsize)
        state.λ = lambda_sample(dm,state.x)
        LambdaU,LambdaV,lambdaa, lambdab = getparams_λ(state.λ,dm)
        U,a,V,b = getparams(state.x,dm);
        grad_sum_U,grad_sum_a,grad_sum_V, grad_sum_b = getparams(state.grad_sum,dm);
        uxx = []
        ixx = []
        for k = 1:state.niters

            state.t += 1
            eps = state.stepsize(state.t)
            err = Array(Float64,dm.batchsize);

            idx, batch = getbatch(dm)
            rr = batch[:,3];
            uu = batch[:,1];
            ii = batch[:,2];

            ux, ix     = unique(uu), unique(ii)
            if k == 1
                uxx = ux
                ixx = ix
            else
                uxx = vcat(uxx,ux)
                ixx = vcat(ixx,ix)
            end
            c_wrappers.c_comp_error!(err,rr,dm.mean_rating,U,V,a,b,uu,ii,dm.batchsize,dm.d);
            halfstep = 0.5 * eps
            sqrtstep = sqrt(eps)

            c_wrappers.c_comp_grad_sum!(err, uu, ii, U, V, a, b, grad_sum_U, grad_sum_V, grad_sum_a, grad_sum_b, dm.batchsize, dm.d)

            c_wrappers.c_update_para_prior!(ux, U, grad_sum_U, LambdaU, zeros(U), zeros(U), dm.inv_pr_pick_users, factor, halfstep, sqrtstep, dm.d)
            c_wrappers.c_update_para_prior!(ix, V, grad_sum_V, LambdaV, zeros(V), zeros(V), dm.inv_pr_pick_items, factor, halfstep, sqrtstep, dm.d)
            c_wrappers.c_update_para_prior!(ux, a, grad_sum_a, lambdaa, zeros(a), zeros(a), dm.inv_pr_pick_users, factor, halfstep, sqrtstep, 1)
            c_wrappers.c_update_para_prior!(ix, b, grad_sum_b, lambdab, zeros(b), zeros(b), dm.inv_pr_pick_items, factor, halfstep, sqrtstep, 1)
        end
        uxx = unique(uxx)
        ixx = unique(ixx)
        uxx, ixx
end

end


#Load some packages and add workers
using Images
using Distributed

using LinearAlgebra
using Statistics

#Load the package on all workers
addprocs(4)
@everywhere using DPMMSubClusters

#Load image, taken from SINTEL http://sintel.is.tue.mpg.de/
img = load("frame_0001.png")

#Change to channel view
img_channels = channelview(img)
z, x, y = size(img_channels)

#Create input
input_arr = zeros(5, x * y)
for i = 1:x
    for j = 1:y
        input_arr[:, (i-1)*y+j] =
            [img_channels[1, i, j], img_channels[2, i, j], img_channels[3, i, j], i, j]
    end
end

#Create HyperParams
#The rgb,xy multiplier allows us to both play with the weight of the xy/rgb
rgb_prior_multiplier = 1
xy_prior_multiplier = 0.1

data_cov = cov(input_arr')
data_cov[4:5, 1:3] .= 0
data_cov[1:3, 4:5] .= 0

data_cov[1:3, 1:3] .*= rgb_prior_multiplier
data_cov[4:5, 4:5] .*= xy_prior_multiplier

data_mean = mean(input_arr, dims = 2)[:]

hyper_params = DPMMSubClusters.niw_hyperparams(1.0, data_mean, 8, data_cov)

#Run the model
labels, clusters, weights =
    DPMMSubClusters.fit(input_arr, hyper_params, 50000.0, iters = 200, verbose = true)

#Get the cluster color means
color_means = [x.μ[1:3] for x in clusters]

segemnated_image = zeros(3, x, y)
for i = 1:x
    for j = 1:y
        segemnated_image[:, i, j] = color_means[labels[(i-1)*y+j]]
    end
end
segemnated_image = colorview(RGB, segemnated_image)


rand(rng::AbstractRNG, ::Random.SamplerType{E}) where {T<:AbstractFloat, E<:RealLog{T}} =
	exp(E, -randexp(rng, T))
rand(rng::AbstractRNG, ::Random.SamplerType{E}) where {E<:RealLog} =
	exp(E, -randexp(rng))

# todo: sample CLogarithmic
# (note that rand(Complex) samples uniformly from the unit square)


"""
Struct for MultiVariate (log moneyness x-> S0*exp(x) ) Copula Process

		gaussianCopulaNVariateLogProcess=GaussianCopulaNVariateLogProcess(models::num1,λ::num2,p::num3,λ₊::num4,λ₋::num5) where {num1,num2,num3,num4,num5 <: Number}

Where:\n
		models  =	the processes.
		rho  = 	correlation matrix.
"""
mutable struct GaussianCopulaNVariateLogProcess{num3 <: Number, numtype <: Number} <: NDimensionalMonteCarloProcess{numtype}
    models::Tuple{Vararg{BaseProcess}}
    rho::Matrix{num3}
    function GaussianCopulaNVariateLogProcess(rho::Matrix{num3}, models::BaseProcess...) where {num3 <: Number}
        sz = size(rho)
        @assert sz[1] == sz[2]
        @assert length(models) == sz[1]
        @assert isposdef(rho)
        zero_typed = predict_output_type_zero_(models...) + zero(num3)
        return new{num3, typeof(zero_typed)}(models, rho)
    end
    function GaussianCopulaNVariateLogProcess(models::BaseProcess...)
        len_ = length(models)
        return GaussianCopulaNVariateLogProcess(Matrix{Float64}(I, len_, len_), models...)
    end
    function GaussianCopulaNVariateLogProcess(model1::BaseProcess, model2::BaseProcess, rho::num3) where {num3 <: Number}
        corr_matrix_ = [1.0 rho; rho 1.0]
        @assert isposdef(corr_matrix_)
        return GaussianCopulaNVariateLogProcess(corr_matrix_, model1, model2)
    end
end

support_type(z::GaussianCopulaNVariateLogProcess{num, num1}) where {num <: Number, num1 <: Number} = zero(num)

export GaussianCopulaNVariateLogProcess;

function simulate!(S_Total, mcProcess::GaussianCopulaNVariateLogProcess, rfCurve::AbstractZeroRateCurve, mcBaseData::AbstractMonteCarloConfiguration, T::Number)
    Nsim = mcBaseData.Nsim
    Nstep = mcBaseData.Nstep
    @assert T > 0.0

    ####Simulation
    ## Simulate
    len_ = length(mcProcess.models)
    #S_Total::Array{Matrix{Number}}=[log.(simulate(model_i,rfCurve,mcBaseData,T)./model_i.underlying.S0) for model_i in mcProcess.models];

    for i = 1:len_
        simulate!(S_Total[i], mcProcess.models[i], rfCurve, mcBaseData, T)
    end

    rho = mcProcess.rho
    if (isdiag(rho))
        return
    end
    for i = 1:len_
        S_Total[i] .= log.(S_Total[i] ./ mcProcess.models[i].underlying.S0)
    end
    U_joint = Matrix{eltype(rho)}(undef, len_, Nsim)
    for j = 1:Nstep
        gausscopulagen2!(U_joint, rho, mcBaseData)
        for i = 1:len_
            @views tmp_ = S_Total[i][:, j+1]
            sort!(tmp_)
            @views tmp_ .= (mcProcess.models[i].underlying.S0) .* exp.(Statistics.quantile(tmp_, U_joint[i, :]; sorted = true))
        end
    end

    ## Conclude
    return
end


mutable struct NeuralNetwork
    layers::AbstractArray{AbstractLayer,1}
    cost_function::AbstractCostFunction
    n_layers::Int64 # cache the number of layers
    function NeuralNetwork(
        layers_::AbstractArray{AbstractLayer,1},
        cost_function_::AbstractCostFunction,
    )
        new(copy(layers_), cost_function_, size(layers_)[1])
    end
end

function fitting(
    nn::NeuralNetwork,
    X_::AbstractArray{Float64,2},
    t::AbstractArray{Float64,2},
    learning_rate::Float64,
    policy = "Adam",
)
    # X = [x_1 x_2 ,,, x_N]
    # t = [t_1 t_2 ,,, t_N]
    X = copy(X_)
    for layer in nn.layers
        X = forward_propagation(layer, X)
    end

    typeof(X) == typeof(t) || error("fitting: size does not match")

    diff = X - t

    for layer in reverse(nn.layers)
        diff = backward_propagation(layer, diff, learning_rate, policy)
    end

end

function predict(nn::NeuralNetwork, X_::AbstractArray{Float64,2})
    # X = [x_1 x_2 ,,, x_N]
    X = copy(X_)
    for layer in nn.layers
        X = forward_propagation(layer, X)
    end

    return X
end


export Arc

mobius(a,b,c,d,z) = (a*z+b)./(c*z+d)
mobius(a,b,c,d,z::Number) = isinf(z)?a/c:(a*z+b)./(c*z+d)
mobiusinv(a,b,c,d,z) = mobius(d,-b,-c,a,z)

mobiusD(a,b,c,d,z) = (a*d-b*c)./(d+c*z).^2
mobiusinvD(a,b,c,d,z) = mobiusD(d,-b,-c,a,z)

immutable Arc{T,V<:Real,TT} <: IntervalDomain{TT}
    center::T
    radius::V
    angles::Tuple{V,V}
    Arc(c,r,a) = new(T(c),V(r),Tuple{V,V}(a))
end

Arc{T<:Number,V<:Real,V1<:Real,V2<:Real}(c::T,r::V,t::Tuple{V1,V2}) =
    Arc{promote_type(T,V,V1,V2),
        promote_type(real(T),V,V1,V2),
        Complex{promote_type(real(T),V,V1,V2)}}(c,r,t)
Arc{T<:Number,V<:Real}(c::Vec{2,T},r::V,t::Tuple{V,V}) =
    Arc{Vec{2,promote_type(T,V)},
        promote_type(real(T),V),
        Vec{2,promote_type(real(T),V)}}(c,r,t)
Arc{T<:Number,V<:Real,V1<:Real,V2<:Real}(c::Vec{2,T},r::V,t::Tuple{V1,V2}) =
    Arc{Vec{2,promote_type(T,V,V1,V2)},
        promote_type(real(T),V,V1,V2),
        Vec{2,promote_type(T,V,V1,V2)}}(c,r,t)

Arc(c::Tuple,r,t) = Arc(Vec(c...),r,t)
Arc(c,r,t0,t1) = Arc(c,r,(t0,t1))

Base.complex{V<:Vec}(a::Arc{V}) = Arc(complex(a.center...),a.radius,a.angles)

isambiguous(d::Arc) =
    isnan(d.center) && isnan(d.radius) && isnan(d.angles[1]) && isnan(d.angles[2])
Base.convert{T<:Number,V<:Real}(::Type{Arc{T,V}},::AnyDomain) =
    Arc{T,V}(NaN,NaN,(NaN,NaN))
Base.convert{IT<:Arc}(::Type{IT},::AnyDomain) =
    Arc(NaN,NaN,(NaN,NaN))

arclength(d::Arc) = d.radius*(d.angles[2]-d.angles[1])

function mobiuspars(z0,r,t0,t1)
    c=exp(im*t0/2)+exp(im*t1/2)
    f=exp(im*t0/2)-exp(im*t1/2)
    d=exp(im*t0/2+im*t1/2)*r
    -c,c*(z0+d),f,f*(d-z0)
end

mobiuspars(a::Arc) = mobiuspars(a.center,a.radius,a.angles...)

for OP in (:mobius,:mobiusinv,:mobiusD,:mobiusinvD)
    @eval $OP(a::Arc,z) = $OP(mobiuspars(a)...,z)
end

tocanonical{T<:Number}(a::Arc{T},x) = real(mobius(a,x))
tocanonicalD{T<:Number}(a::Arc{T},x) = mobiusD(a,x)
fromcanonical{T<:Number,V<:Real,TT<:Complex}(d::Arc{T,V,TT},v::AbstractArray) =
    eltype(d)[fromcanonical(d,vk) for vk in v]
fromcanonical{T<:Number,V<:Real,TT<:Complex}(a::Arc{T,V,TT},x) =
    mobiusinv(a,x)
fromcanonicalD{T<:Number}(a::Arc{T},x) = mobiusinvD(a,x)

tocanonical{V<:Vec}(a::Arc{V},x::Vec) =
    tocanonical(complex(a),complex(x...))
fromcanonical{V<:Vec}(a::Arc{V},x::Number) =
    Vec(reim(fromcanonical(complex(a),x))...)
fromcanonicalD{V<:Vec}(a::Arc{V},x::Number) =
    Vec(reim(fromcanonicalD(complex(a),x))...)

## information

Base.issubset(a::Arc,b::Arc) =
    a.center==b.center && a.radius==b.radius &&
                    (b.angles[1]≤a.angles[1]≤a.angles[2]≤b.angles[2] ||
                            b.angles[1]≥a.angles[1]≥a.angles[2]≥b.angles[2])

# Algebra

for op in (:*,:.*)
    @eval begin
        $op(c::Real,d::Arc) =
            Arc(c*d.center,c*d.radius,(sign(c)*d.angles[1],sign(c)*d.angles[2]))
        $op(d::Arc,c::Real) = $op(c,d)
    end
end

for op in (:+,:-,:.+,:.-)
    @eval begin
        $op(c::Number,d::Arc) = Arc($op(c,d.center),d.radius,d.angles)
        $op(d::Arc,c::Number) = Arc($op(d.center,c),d.radius,d.angles)
    end
end

# allow exp(im*Interval(0,1)) for constructing arc
function Base.exp{CMP<:Complex}(d::Interval{CMP})
    @assert isapprox(real(d.a),0) && isapprox(real(d.b),0)
    Arc(0,1,(imag(d.a),imag(d.b)))
end


@polly function kernel_gesummv(alpha, beta, A, B, tmp, x, y)
    n = size(tmp,1)
    for i = 1:n
        tmp[i] = zero(eltype(tmp))
        y[i] = zero(eltype(tmp))
        for j = 1:n
            tmp[i] = A[i,j] * x[j] + tmp[i]
            y[i] = B[i,j] * x[j] + y[i]
        end
        y[i] = alpha * tmp[i] + beta * y[i]
    end
end

let
    n = 1300

    A = zeros(Float32, n, n)
    B = zeros(Float32, n, n)
    tmp = zeros(Float32, n)
    x = zeros(Float32, n)
    y = zeros(Float32, n)

    alpha = Float32(1.5)
    beta = Float32(1.2)
    for i = 1:n
        x[i] = (i % n) / n
        for j = 1:n
            A[i,j] = ((i*j+1) % n) / n
            B[i,j] = ((i*j+2) % n) / n
        end
    end

    SUITE["gesummv"] = @benchmarkable kernel_gesummv($alpha, $beta, A, B, tmp, x, y) setup = (A = copy($A); B = copy($B); tmp = copy($tmp); x = copy($x); y = copy($y))
end


# OptiGas: Optimize Gas Network using Benders Algorithm
# Author: Edward J. Xu, edxu96@outlook.com
# Date: Aug 12, 2019

using LinearAlgebra
using MatrixOptim

function getData()
    vec_nameNodes = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J"]
    numNodes = length(vec_nameNodes)
    vec_xNode = [1    3    8    2    4    5    6    7    9    9]
    vec_yNode = [1    2    1    7    9    5    3    7    9    4]
    mat_arcTwoNodes = [
        0    1    0    0    0    0    0    0    0    0;
        0    0    0    0    0    1    0    0    0    0;
        0    0    0    0    0    0    0    0    0    1;
        1    0    0    0    1    0    0    0    0    0;
        0    0    0    0    0    1    0    0    0    0;
        1    0    0    1    0    0    0    1    0    0;
        0    1    1    0    0    0    0    0    0    0;
        0    0    0    0    0    0    0    0    1    0;
        0    0    0    0    1    0    0    0    0    0;
        0    0    0    0    0    1    1    0    1    0
    ]
    vec_netEject = [0    0   0    0    0   90    0    0   0   80 ]
    vec_netInject = [10   30  20   10   40   0   25   20  15    0 ]
    mat_distance = zeros(Float64, numNodes, numNodes)
    for m = 1: numNodes
        for n = 1: numNodes
            mat_distance[m, n] = floor(sqrt( (vec_xNode[m] - vec_xNode[n]) * (vec_xNode[m] - vec_xNode[n]) +
                                       (vec_yNode[m] - vec_yNode[n]) * (vec_yNode[m] - vec_yNode[n]) ) )
        end
    end
    mat_fixedCost = zeros(Float64, numNodes, numNodes)
    for m = 1: numNodes
        for n = 1: numNodes
            mat_fixedCost[m, n] = 10 * mat_distance[m, n]
        end
    end
    return numNodes, vec_netEject, vec_netInject, mat_distance, mat_fixedCost
end

function getModelQ4(numNodes, vec_netEject, vec_netInject, mat_distance, mat_fixedCost)
    # Fomulation of mat_aCap
    mat_aCap_1 = zeros(numNodes^2, numNodes^2)
    for mn = 1: numNodes^2
        mat_aCap_1[mn, mn] = -1
    end
    # mat_a2_1 = zeros(numNodes^2, numNodes^2)
    mat_aCap_3 = zeros(numNodes, numNodes^2)
    seq_zeroTenNinety = collect(0:10:90)
    for m = 1: numNodes
        mat_aCap_3[m, (10 * m - 9): (10 * m)] = repeat([-1], 10)  # sent out
        mat_aCap_3[m, (seq_zeroTenNinety + repeat([m],10))] = repeat([1], 10)  # sent in
        mat_aCap_3[:, (11 * m - 10)] = repeat([0], 10)  # self-sending doesn't count
    end
    mat_aCap = vcat(mat_aCap_1, mat_aCap_3)

    # Fomulation of Matrix B
    mat_bCap_1 = zeros(numNodes^2, numNodes^2)
    for mn = 1: numNodes^2
        mat_bCap_1[mn, mn] = 170
    end
    # mat_b2_1 = Diagnal(repeat([1], numNodes^2))
    mat_bCap_3 = zeros(numNodes, numNodes^2)
    mat_bCap = vcat(mat_bCap_1, mat_bCap_3)

    # Fomulation of Vector b
    vec_bCap_1 = zeros(numNodes^2)
    # vec_b2_1 = zeros(numNodes^2)
    # for m = 1: numNodes
    #     for n = 1: numNodes
    #         vec_b2_1[10 * (m - 1) + n] = mat_arcTwoNodes[m, n]
    #     end
    # end
    vec_bCap_3 = zeros(numNodes)
    for n = 1: numNodes
        vec_bCap_3[n] = vec_netInject[n] - vec_netEject[n]
    end
    vec_b = vcat(vec_bCap_1, vec_bCap_3)
    vec_b = hcat(vec_b)

    vec_y_max = repeat([1], numNodes^2)
    vec_y_max = hcat(vec_y_max)
    vec_y_min = zeros(numNodes^2)
    for m = 1: numNodes
        for n = 1: numNodes
            vec_y_min[10 * (m - 1) + n] = mat_arcTwoNodes[m, n]
        end
    end

    vec_c = zeros(numNodes^2)
    for m = 1: numNodes
        for n = 1: numNodes
            vec_c[10 * (m - 1) + n] = mat_distance[m, n]
        end
    end
    vec_c = hcat(vec_c)

    vec_f = zeros(numNodes^2)
    for m = 1: numNodes
        for n = 1: numNodes
            vec_f[10 * (m - 1) + n] = mat_fixedCost[m, n]
        end
    end
    vec_f = hcat(vec_f)

    return vec_c, mat_aCap, vec_b, vec_f, vec_bCap, vec_y_min, vec_y_max
end

function main()
    numNodes, vec_netEject, vec_netInject, mat_distance, mat_fixedCost = getData()
    vec_c, mat_aCap, vec_b, vec_f, vec_bCap, vec_y_min, vec_y_max = getModelQ4(
        numNodes, vec_netEject, vec_netInject, mat_distance, mat_fixedCost
    )
    solveBendersMilp(
        n_x = numNodes^2,
        n_y = numNodes^2,
        vec_min_y = vec_y_min,
        vec_max_y = vec_y_max,
        vec_c = vec_c,
        vec_f = vec_f,
        vec_b = vec_b,
        mat_a = mat_aCap,
        mat_b = mat_bCap,
        epsilon = 0.0001,
        timesIterationMax = 1000
    )
    # 1280 - sum(mat_arcTwoNodes .* mat_fixedCost)
end

main()


using StatsBase

#NETWORK PROPERTIES
LAYER_SIZE = 3
SEQUENCE_LEN = 25
LEARN_RATE = 1e-1

#NETWORK WEIGHTS, BIASES
function weights(vsize)
	W_XH = randn(LAYER_SIZE,vsize)*0.01
	W_HH = randn(LAYER_SIZE,LAYER_SIZE)*0.01
	W_HY = randn(vsize,LAYER_SIZE)*0.01
	return W_XH, W_HH, W_HY
end

function biases(vsize)
	B_H = zeros(LAYER_SIZE,1)
	B_Y = zeros(vsize,1)
	return B_H, B_Y
end

#LOSS FUNCTION

function loss(inputs,targets,vsize,W_XH,W_HH,W_HY,B_H,B_Y,LAST_H)
	xs, hs, ys, ps = Dict(), Dict(), Dict(), Dict()
	hs[0] = copy(LAST_H)
	error = 0

	for t in 1:length(inputs)
		xs[t] = zeros(vsize,1)
		xs[t][inputs[t]] = 1
		hs[t] = tanh((W_XH*xs[t]) + (W_HH*hs[t-1]) + B_H)
		ys[t] = (W_HY*hs[t]) + B_Y
		ps[t] = exp(ys[t])/sum(exp(ys[t]))
		error += -log(ps[t][targets[t],1])
	end

	D_W_XH, D_W_HH, D_W_HY = zeros(W_XH), zeros(W_HH), zeros(W_HY)
	D_B_H, D_B_Y = zeros(B_H), zeros(B_Y)
	H_NEXT = zeros(hs[0])

	for t in reverse(1:length(inputs))
		dy = copy(ps[t])
		dy[targets[t]] -= 1
		D_W_HY += dy*hs[t].'
		D_B_Y += dy
		dh = (W_HY.'*dy) + H_NEXT
		dh_raw = (1 - hs[t].*hs[t]).*dh
		D_B_H += dh_raw
		D_W_XH += dh_raw*xs[t].'
		D_W_HH += dh_raw*hs[t-1].'
		H_NEXT = W_HH'*dh_raw
	end

	D_W_XH = clamp!(D_W_XH,-5,5)
	D_W_HH = clamp!(D_W_HH,-5,5)
	D_W_HY = clamp!(D_W_HY,-5,5)
	D_B_H = clamp!(D_B_H,-5,5)
	D_B_Y = clamp!(D_B_Y,-5,5)

	return error, D_W_XH, D_W_HH, D_W_HY, D_B_H, D_B_Y, hs[length(inputs)]

end

#SAMPLER

function sample(h, seed, vsize, n, W_XH, W_HH, W_HY, B_H, B_Y)
	x = zeros(vsize,1)
	x[seed] = 1
	ixes = []

	for t in 1:n
		h = tanh(W_XH*x + W_HH*h + B_H)
		y = W_HY*h + B_Y
		p = exp(y)/sum(exp(y))
		w = StatsBase.weights(p)
		ix = StatsBase.sample(1:vsize,w)
		x = zeros(vsize,1)
		x[ix] = 1
		append!(ixes,ix)
	end

	return ixes
end

#WRITE TO FILE

function tofile(n, text)
	f = open("output.txt","a")
		write(f,"----------\n$n\n----------\n\n$text\n\n")
	close(f)
end

#INITIALIZE THE NETWORK
function run(contents)
	n,p = 1,1
	chars = Set([ch for ch in contents])
	vsize = length(chars)
	ix_char = Dict(ch for ch in enumerate(chars))
	char_ix = map(reverse,ix_char)
	errs = []

	smooth_loss = -log(1.0/vsize) * SEQUENCE_LEN

	W_XH, W_HH, W_HY = weights(vsize)
	B_H, B_Y = biases(vsize)

	mXh, mHh, mHy, mBh, mBy = zeros(W_XH), zeros(W_HH), zeros(W_HY), zeros(B_H), zeros(B_Y)

	while true
		if p + SEQUENCE_LEN + 1 >= length(contents) || n == 1
			LAST_H = zeros(LAYER_SIZE,1)
			p = 1
		end

		inputs = [char_ix[ch] for ch in contents[p:p+SEQUENCE_LEN]]
		targets = [char_ix[ch] for ch in contents[p+1:p+1+SEQUENCE_LEN]]

		error, D_W_XH, D_W_HH, D_W_HY, D_B_H, D_B_Y, LAST_H = loss(inputs, targets, vsize, W_XH, W_HH, W_HY, B_H, B_Y, LAST_H)

		append!(errs,error)

		if n%100 == 0
			println("$n: ", sum(errs)/length(errs))
		end

		if n%1000 == 0
			text = sample(LAST_H,inputs[1],vsize,200,W_XH,W_HH,W_HY,B_H,B_Y)
			tofile(n, join([ix_char[ch] for ch in text]))
		end

		#MANUAL WEIGHTING
		#XH
		mXh += D_W_XH.*D_W_XH
		W_XH += -LEARN_RATE * D_W_XH ./ sqrt(mXh + 1e-8)
		#HH
		mHh += D_W_HH.*D_W_HH
		W_HH += -LEARN_RATE * D_W_HH ./ sqrt(mHh + 1e-8)
		#HY
		mHy += D_W_HY.*D_W_HY
		W_HY += -LEARN_RATE * D_W_HY ./ sqrt(mHy + 1e-8)
		#BH
		mBh += D_B_H.*D_B_H
		B_H += -LEARN_RATE * D_B_H ./ sqrt(mBh + 1e-8)
		#BY
		mBy += D_B_Y.*D_B_Y
		B_Y += -LEARN_RATE * D_B_Y ./ sqrt(mBy + 1e-8)

		p += SEQUENCE_LEN
		n += 1
	end

end

#LOAD FILE

f = open("texts/gatsby_ord.txt")
	contents = readstring(f)
close(f)

y = run(contents)


using DSP
using WAV
function AD_chebyshev(order,x)

    return-(x*cos(order*acos(x))+order*sqrt(1-x^2)*sin(order*acos(x)))/(-1+order^2)

end
#=
this code is based on this post
https://jatinchowdhury18.medium.com/practical-considerations-for-antiderivative-anti-aliasing-d5847167f510

this is very similar reworking of his code for my purposes
=#

#anti derivative anti aliasing order 1
function processAD1(signal::Vector{Float64},f,f1, tolerance = 1.0e-5)

    x1 = 0.0
    returnSignal = Float64[]
    for s in signal
        if abs(s - x1) < tolerance
            push!(returnSignal, f((s + x1) / 2))
        else
            push!(returnSignal, (f1(s) - f1(x1)) / (s-x1))
            x1 = s
        end
    end
    return returnSignal
end
#anti derivative anti aliasing order 2
function processAD2(signal, f, f1, f2, ε=1.0e-5)

    returnSignal = Float64[]

    calcD(x0, x1) = abs(x0-x1) < ε ? f1((x0+x1)/2) : (f2(x0) - f2(x1)) / (x0-x1)

    x₁ = 0.0
    x₂ = 0.0
    for x₀ in signal
        if abs(x₀ - x₁) < ε
            x̄ = (x₀ + x₂)/2
            Δ = x̄ - x₀
            if Δ < ε
                push!(returnSignal,  f((x̄  + x₀) / 2.0))
            else
                push!(returnSignal,  (2/Δ) * (f1(x̄) + (f2(x₀) - f2(x̄)) / Δ))
            end
        else
           push!(returnSignal, (2.0 / (x₀ - x₂)) * (calcD(x₀, x₁) - calcD(x₁, x₂)) )
           x₂ = x₁
           x₁ = x₀
        end
    end

    return returnSignal
end

#hardclip first antiderivatie
function hc1(x)
    return x*sign(x) - .5
end
#hardclip second order derivative
function hc2(x)
    return (x^2/2 + 1/6) * sign(x) - x/2
end

function sc(x::Float64)::Float64
    return 2*atan(x)/pi
end
#softClip first order anti derivative
function sc1(x::Float64)::Float64
    return (2x*atan(x)-log(x^2+1))/pi
end
#softClip second order anti derivatie
function sc2(x::Float64)::Float64
    return (x*(-log(x^2+1))+(x^2-1)*(atan(x))+x)/pi
end
function hardC(signal::Vector{Float64}, gain)::Vector{Float64}
    c(x,y) = (abs(x*y) < 1 && return x*y)  || return sign(x)
    return c.(signal,gain)
end
function hardClipAD1(signal::Vector{Float64}, gain)::Vector{Float64}
    f(x) = (abs(x) < 1 ? x : sign(x) )
    f1(x) = (abs(x) < 1 ? x^2/2 : x*sign(x) - .5)
    return processAD1(gain*signal, f,f1)
end

function hardClipAD2(signal::Vector{Float64}, gain)::Vector{Float64}
    f(x) = (abs(x) < 1 ? x : sign(x) )
    f1(x) = (abs(x) < 1 ? x^2/2 : x*sign(x) - .5)
    f2(x) = (abs(x) < 1 ? x^3/6 : (x^2/2 + 1/6) * sign(x) - x/2)
    return processAD2(gain*signal, f,f1,f2)
end
function softClipAD2(signal::Vector{Float64}, gain)::Vector{Float64}
    return processAD2(gain*signal, sc,sc1,sc2)
end


using Learn
using Test

@testset "Learn.jl" begin
    # Write your tests here.
    @testset "Add" begin
        @test true
        @test Learn.add(3,5) == 8
        @test Learn.add(1.1, 3.0) == 4.1
        @test Learn.add(1+5im, 4-2im) == 5+3im
        @test Learn.add(Vector([1,2,3]), Vector([4,5,6])) == Vector([5,7,9])
        @test Learn.add("asdf ", "fdsa") == "asdf fdsa"
    end
    @testset "Sum" begin
        @test Learn.∑(1,2,3,4,5) == 15
        @test Learn.∑([1,2,3,4,5]) == 15
    end

    include("testDayCounts.jl")
    include("testHolidays.jl")
end


# test default solvers

include("linprog.jl")
linprogtest()

include("mixintprog.jl")
mixintprogtest()

include("quadprog.jl")
quadprogtest()

# Test conic fallback for LPs
include("conicinterface.jl")
coniclineartest(MathProgBase.defaultLPsolver, duals=true)

# Test LP fallback for conics
include("linproginterface.jl")
linprogsolvertest(MathProgBase.defaultConicsolver)


"Mixture of Factor Analyzers"
function fit_mm(::Type{FactorAnalysis}, X::AbstractMatrix{T},
                k::Integer;            # the number of factor analyzers to use
                factors::Integer = 1,  # the number of factors in each analyzer
                tol::Real=1.0e-6,      # convergence tolerance
                maxiter::Integer=1000, # number of iterations
                μs::Union{AbstractMatrix{T}, Nothing} = nothing,
                Σs::Union{AbstractArray{T,3}, Nothing} = nothing
            ) where {T<:AbstractFloat}

    d, n = size(X)
    CMVN = (2π)^(-d/2)
    mX=mean(X, dims=2)
    cX=cov(X, dims=2)
    sc=det(cX)^(1/d)
    l = factors

    # initialize parameters
    πⱼ = fill(one(T)/k, k)
    Wⱼ = zeros(T,d,l,k)
    μⱼ = zeros(T,d,k)
    Ψⱼ = zeros(T,d,k)
    rn = 1:min(d,l)
    for j in 1:k
        if Σs !== nothing
            ev = eigvecs(Σs[:,:,j], sortby=-)
            Wⱼ[:,rn,j] = ev[:,rn].*sqrt(det(Σs[:,:,j])^(1/d))
        else
            Wⱼ[:,:,j] = randn(d,l) * sqrt(sc/l)
        end
        μⱼ[:,j] = μs !== nothing ? μs[:,j] : (vec(randn(d)' * sqrt(cX)) .+ mX)
        Ψⱼ[:,j] = diag(Σs !== nothing ? Σs[:,:,j] : cX) .+ eps()
    end

    hᵢⱼ = zeros(T,k,n)
    πᵢⱼ = similar(hᵢⱼ)
    E₍zzᵀₗxᵢωⱼ₎ = zeros(T,l,l,k)
    Σ⁻¹ⱼ = zeros(T,d,d,k)
    ZIᵀ=zeros(T,n,l+1,k)
    ZIᵀ[:,end,:] .= one(T)
    Σ⁻¹Y = similar(X)
    Y = similar(X)

    E₍zₗxᵢωⱼ₎ = view(ZIᵀ,:,1:l,:)

    L_old = 0.0
    for itr in 1:maxiter
        # E Step
        # 1) Calculate hᵢⱼ, E₍zₗxᵢωⱼ₎
        for j in 1:k
            Ψ⁻¹ = Diagonal(1 ./ Ψⱼ[:, j])
            W = @view Wⱼ[:,:,j]
            Σ⁻¹ = @view Σ⁻¹ⱼ[:,:,j]
            # V⁻¹ = inv(I + WᵀΨ⁻¹*W)
            WᵀΨ⁻¹ = W'*Ψ⁻¹
            # Σ⁻¹ = Ψ⁻¹ - Ψ⁻¹*W*V⁻¹*WᵀΨ⁻¹
            Σ⁻¹[:] = Ψ⁻¹ - Ψ⁻¹*W*inv(I + WᵀΨ⁻¹*W)*WᵀΨ⁻¹
            detΣ⁻¹ = sqrt(det(Σ⁻¹))
            # Y = X .- view(μⱼ,:,j)
            broadcast!(-,Y,X,view(μⱼ,:,j))
            Σ⁻¹Y .= Σ⁻¹*Y
            # view(hᵢⱼ,j,:) .= πⱼ[j]*(CMVN*detΣ⁻¹).*exp.(-0.5.*vec(sum(Y.*Σ⁻¹Y, dims=1)))
            broadcast!(*,Y,Y,Σ⁻¹Y)
            sum!(view(hᵢⱼ,j,:)', Y)
            c = πⱼ[j]*(CMVN*detΣ⁻¹)
            broadcast!(e->c*exp(-0.5*e),view(hᵢⱼ,j,:),view(hᵢⱼ,j,:))
            # Ez[:,:,j] = V⁻¹*WᵀΨ⁻¹*Y # W'*Σ⁻¹*Y
            view(E₍zₗxᵢωⱼ₎,:,:,j) .= (W'Σ⁻¹Y)'
        end

        # 2) Calculate log(L)
        hᵢⱼ[hᵢⱼ .< eps()] .= eps()
        hᵢ = sum(hᵢⱼ, dims=1)
        L = sum(log.(hᵢ))

        # 3) Calculate πⱼ
        # πᵢⱼ = hᵢⱼ./hᵢ
        broadcast!(/,πᵢⱼ,hᵢⱼ,hᵢ)
        πⱼ .= sum(πᵢⱼ, dims=2) |> vec

        # Check exit conditions
        if itr > 1
            chg = L - L_old
            if chg < 0
                @warn "log likelihood decreased" itr=itr change=chg
            end
            @debug "log(L)=$L" itr=itr increment=chg πⱼ=πⱼ' # Ψⱼ=Ψⱼ
            if (abs(chg) < tol)
                break
            end
        end
        L_old = L

        # 4) Calculate E₍zzᵀₗxᵢωⱼ₎
        for j in 1:k
            W = @view Wⱼ[:,:,j]
            Σ⁻¹ = @view Σ⁻¹ⱼ[:,:,j]
            broadcast!(-,Y,X,view(μⱼ,:,j))
            βⱼ = W'*Σ⁻¹
            Sᵢ = (view(πᵢⱼ, j, :)'.*Y)*Y'./πⱼ[j]
            view(E₍zzᵀₗxᵢωⱼ₎,:,:,j) .= I - βⱼ*W + βⱼ*Sᵢ*transpose(βⱼ)
        end

        # M Step
        for j in 1:k
            πᵢ = @view πᵢⱼ[j, :]
            broadcast!(*,Y,πᵢ',X) # πᵢ'.*X
            QWnum = Y*view(ZIᵀ,:,:,j)
            QWden1 = πⱼ[j]*view(E₍zzᵀₗxᵢωⱼ₎,:,:,j)
            QWden23 = view(E₍zₗxᵢωⱼ₎,:,:,j)'*πᵢ
            QWden = [QWden1 QWden23; QWden23' πⱼ[j]]
            R = QWnum/QWden

            Wⱼ[:,:,j] .= R[:,1:end-1]
            μⱼ[:,j] .= R[:,end]

            Ψⱼ[:,j] .= diag(Y*X' - R*QWnum')./πⱼ[j]
        end
    end
    return FinateMixtureModel([FactorAnalysis(μⱼ[:,j], Wⱼ[:,:,j], Ψⱼ[:,j]) for j in 1:k], πⱼ./n)
end

fit_mm(::Type{FactorAnalysis}, X::AbstractMatrix{T}; m=2, k=1, kwargs...) where {T<:AbstractFloat} =
    fit_mm(FactorAnalysis, X, m; factors=k, kwargs...)