##############################################################################
##
## Non Linear solver using Pseudo-Transient Continuation Method
##
##############################################################################

# Implicit time step
function implicit_timestep(G!, ypost, Δ; is_algebraic = fill(false, size(ypost)...), iterations = 100, verbose = true, method = :newton, autodiff = :forward, maxdist = sqrt(eps()), J0c = (nothing, nothing), y̲ = fill(-Inf, length(ypost)), ȳ = fill(Inf, length(ypost)), reformulation = :smooth)
    G_helper!(ydot, y) = (G!(ydot, y) ; ydot .-= .!is_algebraic .* (ypost .- y) ./ Δ)
    J0, colorvec = J0c
    if J0 == nothing
        result = nlsolve(G_helper!, ypost; iterations = iterations, show_trace = verbose, ftol = maxdist, method = method, autodiff = autodiff)
    else
        if autodiff == :forward
            jac_cache = ForwardColorJacCache(G_helper!, deepcopy(ypost); colorvec = colorvec, sparsity = J0)
            j_helper! = (J, y) -> forwarddiff_color_jacobian!(J, G_helper!, y, jac_cache)
        else
            j_helper! = (J, y) -> finite_difference_jacobian!(J, G_helper!, y; colorvec = colorvec)
        end
        if any(y̲ .!= -Inf) || any(ȳ .!= Inf)
            # using mcpsolve if lower/upper bounds are given
            result = mcpsolve(OnceDifferentiable(G_helper!, j_helper!, deepcopy(ypost), deepcopy(ypost), J0), y̲, ȳ, ypost; iterations = iterations, show_trace = verbose, ftol = maxdist, method = method, reformulation = reformulation)
        else
            result = nlsolve(OnceDifferentiable(G_helper!, j_helper!, deepcopy(ypost), deepcopy(ypost), J0), ypost; iterations = iterations, show_trace = verbose, ftol = maxdist, method = method)
        end
    end
    return result.zero, result.residual_norm
end

# Solve for steady state
function finiteschemesolve(G!, y0; Δ = 1.0, is_algebraic = fill(false, size(y0)...), iterations = 100, inner_iterations = 10, verbose = true, inner_verbose = false, method = :newton, autodiff = :forward, maxdist = sqrt(eps()), scale = 10.0, J0c = (nothing, nothing), minΔ = 1e-9, y̲ = fill(-Inf, length(y0)), ȳ = fill(Inf, length(y0)), reformulation = :smooth)
    ypost = y0
    ydot = zero(y0)
    G!(ydot, ypost)
    residual_norm = norm(ydot) / length(ydot)
    isnan(residual_norm) && throw("G! returns NaN with the initial value")
    if Δ == Inf
        ypost, residual_norm = implicit_timestep(G!, y0, Δ; is_algebraic = is_algebraic, verbose = verbose, iterations = iterations,  method = method, autodiff = autodiff, maxdist = maxdist, J0c = J0c, y̲ = y̲, ȳ = ȳ)
    else
        coef = 1.0
        oldresidual_norm = residual_norm
        iter = 0
        if verbose
            @printf "Iter   TimeStep   Residual\n"
            @printf "---- ---------- ----------\n"
        end
        while (iter < iterations) & (Δ >= minΔ) & (residual_norm > maxdist)
            iter += 1
            y, nlresidual_norm = implicit_timestep(G!, ypost, Δ; is_algebraic = is_algebraic, verbose = inner_verbose, iterations = inner_iterations, method = method, autodiff = autodiff, maxdist = maxdist, J0c = J0c, y̲ = y̲, ȳ = ȳ, reformulation = reformulation)
            G!(ydot, y)
            if any(y̲ .!= -Inf) || any(ȳ .!= Inf)
                mask = y̲ .+ eps() .<= y .<= ȳ .- eps() # only unconstrained ydot is relevant for residual_norm calculation
                residual_norm, oldresidual_norm = norm(ydot .* (mask))/sum(mask), residual_norm
            else
                residual_norm, oldresidual_norm = norm(ydot) / length(ydot), residual_norm
            end
            residual_norm = isnan(residual_norm) ? Inf : residual_norm
            if nlresidual_norm <= maxdist
                # if the implicit time step is correctly solved
                if verbose
                    @printf "%4d %8.4e %8.4e\n" iter Δ residual_norm
                end
                coef = (residual_norm <= oldresidual_norm) ? scale * coef : 1.0
                Δ = Δ * coef * oldresidual_norm / residual_norm
                ypost, y = y, ypost
            else
                # verbose && @show iter, Δ, NaN
                # if the implict time step is not solved
                # revert and diminish the time step
                coef = 1.0
                Δ = Δ / 10
                residual_norm = oldresidual_norm
            end
        end
    end
    verbose && ((residual_norm > maxdist) | (Δ < minΔ)) && @warn "Iteration did not converge"
    return ypost, residual_norm
end


# mercer mapping

export MercerMap, apply, linear, rbf

immutable MercerMap
    X::Matrix           # original reference data
    kernel::Symbol      # kernel function
    K::Matrix           # (centered) kernel matrix of reference data
    d::Vector           # largest eigenvalues of kernel matrix
    V::Matrix           # eigenvectors of kernel matrix
end

# constructing mercer map
function MercerMap(X::Matrix, kernel::Symbol)
    kernelfcn = eval(kernel)
    K = center(kernelfcn(X, X))

    d, V = eig(Symmetric(K))

    # consider dimensions with eigenvalues > 1e-9
    i = d .> 1e-9
    V = V[:,i]
    d = sqrt(d[i])

    MercerMap(X, kernel, K, d, V)
end

# map reference data
apply(map::MercerMap) = map.V*diagm(map.d)
function apply(map::MercerMap, X::Matrix)
    kernelfcn = eval(map.kernel)
    KT = center(map.K, kernelfcn(map.X, X))
    KT'*map.V*diagm(1./map.d)
end

# center data to origin in feature space (M. Meila: Data Centering in Feature Space, Eq. 17)
function center(K::Matrix, KT::Matrix)
    m = size(K,1)
    M1 = 1/m*ones(size(K))
    N1 = ones(size(KT)) / m
    KT = KT - K*N1 - M1*KT + M1*K*N1
end
center(K::Matrix) = center(K, K)

# kernel functions
linear(X::Matrix, Y::Matrix) = full(X*Y')
rbf(X::Matrix, Y::Matrix; sigma::Float64=1.0) = full(exp(-dist2(X, Y)/(2*sigma^2)))

function dist2(X::Matrix, Y::Matrix)
    xx = sum(X.*X, 2)
    yy = sum(Y.*Y, 2)
    xy = X*Y'
    abs(repmat(xx, 1, size(yy, 1)) + repmat(yy', size(xx, 1), 1) - 2*xy)
end


# using SymPy

# x, y = @vars x y
# u = x*(1-x)*y*(1-y)
# v = x*(1-x)*y*(1-y)^2

# p = x*(1-x)*y*(1-y)

# ux = diff(u,x)
# uy = diff(u,y)
# vx = diff(v,x)
# vy = diff(v,y)
# px = diff(p,x)
# py = diff(p,y)
# f = u*ux + v*uy - (diff(ux,x)+diff(uy,y)) + px
# g = u*vx + v*vy - (diff(vx,x)+diff(vy,y)) + py
# h = diff(u, x) + diff(v, y)
# println(replace(replace(sympy.julia_code(simplify(f)), ".^"=>"^"), ".*"=>"*"))
# println(replace(replace(sympy.julia_code(simplify(g)), ".^"=>"^"), ".*"=>"*"))
# println(replace(replace(sympy.julia_code(simplify(h)), ".^"=>"^"), ".*"=>"*"))

function u_exact(x,y)
    if y == 0.0
        return 1.0
    else
        return 0.0
    end
end

function v_exact(x,y)
    0.0
end

function p_exact(x,y)
    0.0
end

function ffunc(x, y)
    #-x^2*y*(x - 1)^2*(y - 1)^2*(2*y - 1) + x*y^2*(x - 1)*(2*x - 1)*(y - 1)^2 + x*y*(y - 1) - 2*x*(x - 1) + y*(x - 1)*(y - 1) - 2*y*(y - 1)
    0.0
end
function gfunc(x, y)
    #x^2*y*(x - 1)^2*(y - 1)^3*(3*y - 1) - x*y^2*(x - 1)*(2*x - 1)*(y - 1)^3 + 3*x*y*(x - 1) + 5*x*(x - 1)*(y - 1) + 2*y*(y - 1)^2
    0.0
end

function hfunc(x,y)
    # (y - 1)*(-2*x*y*(x - 1) + x*y - x*(x - 1)*(y - 1) + y*(x - 1))
    0.0
end

# nn to estimate space varying viscosity
function nu_exact(x, y)
    1 + 6 * x^2 + x / (1 + 2 * y^2)
    # (1 + x^2 + x / (1 + 2 * y^2)) * 0.01
    # 1 + x^2
end

using LinearAlgebra
using MAT
using AdFem
using PyPlot
using SparseArrays

m = 20
n = 20
h = 1/n
# nu = 0.01

F1 = compute_fem_source_term1(eval_f_on_gauss_pts(ffunc, m, n, h), m, n, h)
F2 = compute_fem_source_term1(eval_f_on_gauss_pts(gfunc, m, n, h), m, n, h)
H = h^2*eval_f_on_fvm_pts(hfunc, m, n, h)
B = constant(compute_interaction_matrix(m, n, h))

# compute F
xy = fem_nodes(m, n, h)
x, y = xy[:,1], xy[:,2]
u0 = @. u_exact(x,y)
v0 = @. v_exact(x,y)

xy = fvm_nodes(m, n, h)
x, y = xy[:,1], xy[:,2]
p0 = @. p_exact(x,y)
nu_fvm = constant(@. nu_exact(x, y))
nu_gauss = reshape(repeat(nu_fvm, 1, 4), (-1,))

Laplace = compute_fem_laplace_matrix1(nu_gauss, m, n, h)

# Laplace = constant(nu*compute_fem_laplace_matrix1(m, n, h))
function compute_residual(S)
    u, v, p = S[1:(m+1)*(n+1)], S[(m+1)*(n+1)+1:2(m+1)*(n+1)], S[2(m+1)*(n+1)+1:2(m+1)*(n+1)+m*n]
    G = eval_grad_on_gauss_pts([u;v], m, n, h)
    ugauss = fem_to_gauss_points(u, m, n, h)
    vgauss = fem_to_gauss_points(v, m, n, h)
    ux, uy, vx, vy = G[:,1,1], G[:,1,2], G[:,2,1], G[:,2,2]

    interaction = compute_interaction_term(p, m, n, h) # julia kernel needed
    f1 = compute_fem_source_term1(ugauss.*ux, m, n, h)
    f2 = compute_fem_source_term1(vgauss.*uy, m, n, h)
    f3 = -interaction[1:(m+1)*(n+1)]
    f4 = Laplace*u
    f5 = -F1
    F = f1 + f2 + f3 + f4 + f5

    g1 = compute_fem_source_term1(ugauss.*vx, m, n, h)
    g2 = compute_fem_source_term1(vgauss.*vy, m, n, h)
    g3 = -interaction[(m+1)*(n+1)+1:end]
    g4 = Laplace*v
    g5 = -F2
    G = g1 + g2 + g3 + g4 + g5

    H0 = -B * [u;v] + H

    R = [F;G;H0]
    return R
end

function compute_jacobian(S)
    u, v, p = S[1:(m+1)*(n+1)], S[(m+1)*(n+1)+1:2(m+1)*(n+1)], S[2(m+1)*(n+1)+1:2(m+1)*(n+1)+m*n]
    G = eval_grad_on_gauss_pts([u;v], m, n, h)
    ugauss = fem_to_gauss_points(u, m, n, h)
    vgauss = fem_to_gauss_points(v, m, n, h)
    ux, uy, vx, vy = G[:,1,1], G[:,1,2], G[:,2,1], G[:,2,2]

    M1 = constant(compute_fem_mass_matrix1(ux, m, n, h))
    M2 = constant(compute_fem_advection_matrix1(constant(ugauss), constant(vgauss), m, n, h)) # a julia kernel needed
    M3 = Laplace
    Fu = M1 + M2 + M3

    Fv = constant(compute_fem_mass_matrix1(uy, m, n, h))

    N1 = constant(compute_fem_mass_matrix1(vy, m, n, h))
    N2 = constant(compute_fem_advection_matrix1(constant(ugauss), constant(vgauss), m, n, h))
    N3 = Laplace
    Gv = N1 + N2 + N3

    Gu = constant(compute_fem_mass_matrix1(vx, m, n, h))

    J0 = [Fu Fv
          Gu Gv]
    J = [J0 -B'
        -B spdiag(zeros(size(B,1)))]
end

NT = 5
# S = zeros(m*n+2(m+1)*(n+1), NT+1)
# S[1:m+1, 1] = ones(m+1,1)

bd = bcnode("all", m, n, h)
# bd = [bd; bd .+ (m+1)*(n+1); ((1:m) .+ 2(m+1)*(n+1))]
bd = [bd; bd .+ (m+1)*(n+1); 2*(m+1)*(n+1)+(n-1)*m+1:2*(m+1)*(n+1)+(n-1)*m+2] # only apply Dirichlet to velocity

function solve_steady_cavityflow_one_step(S)
    residual = compute_residual(S)
    J = compute_jacobian(S)

    J, _ = fem_impose_Dirichlet_boundary_condition1(J, bd, m, n, h)
    # residual[bd] .= 0.0
    residual = scatter_update(residual, bd, zeros(length(bd)))

    d = J\residual
    S_new = S - d
    return S_new
end

function condition(i, S_arr)
    i <= NT + 1
end

function body(i, S_arr)
    S = read(S_arr, i-1)
    op = tf.print("i=",i)
    i = bind(i, op)
    S_new = solve_steady_cavityflow_one_step(S)
    S_arr = write(S_arr, i, S_new)
    return i+1, S_arr
end

# for i = 1:NT
#     residual = compute_residual(S[:,i])
#     J = compute_jacobian(S[:,i])

#     J, _ = fem_impose_Dirichlet_boundary_condition1(J, bd, m, n, h)
#     residual[bd] .= 0.0

#     d = J\residual
#     S[:,i+1] = S[:,i] - d
#     @info i, norm(residual)
# end

S_arr = TensorArray(NT+1)
S_arr = write(S_arr, 1, [u0; v0; p0])

i = constant(2, dtype=Int32)

_, S = while_loop(condition, body, [i, S_arr])
S = set_shape(stack(S), (NT+1, 2*(m+1)*(n+1)+m*n))

sess = Session(); init(sess)
output = run(sess, S)

S = output
# out_v = output[:, 1:2*(m+1)*(n+1)]
# out_p = output[:, 2*(m+1)*(n+1)+1:end]

matwrite("steady_cavity_data.mat",
    Dict(
        # "V"=>output[end, 1:2*(m+1)*(n+1)]
        # "V"=>output[:, 1:2*(m+1)*(n+1)]
        "V"=>output
    ))

# figure(figsize=(20,10))
# subplot(321)
# visualize_scalar_on_fem_points(output[NT+1, 1:(m+1)*(n+1)], m, n, h)
# subplot(322)
# visualize_scalar_on_fem_points(u0, m, n, h)

# subplot(323)
# visualize_scalar_on_fem_points(output[NT+1, (m+1)*(n+1)+1:2*(m+1)*(n+1)], m, n, h)
# subplot(324)
# visualize_scalar_on_fem_points(v0, m, n, h)

# subplot(325)
# visualize_scalar_on_fvm_points(output[NT+1, 2*(m+1)*(n+1)+1:end], m, n, h)
# subplot(326)
# visualize_scalar_on_fvm_points(p0, m, n, h)

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

# prev_data = matread("steady_cavity_data.mat")["V"]

# final_u2=prev_data[NT+1, 1:(1+m)*(1+n)]
# final_v2=prev_data[NT+1, (1+m)*(1+n)+1:2*(m+1)*(n+1)]
# final_p2=prev_data[NT+1, 2*(m+1)*(n+1)+1:end]

# u12 = final_u2[Int(n/2)*(m+1)+1: Int(n/2)*(m+1)+m+1]
# u22 = final_u2[Int(n/2)+1:m+1:end]

# v12 = final_v2[Int(n/2)*(m+1)+1: Int(n/2)*(m+1)+m+1]
# v22 = final_v2[Int(n/2)+1:m+1:end]
####################################################################################

# final_u=output[NT+1, 1:(1+m)*(1+n)]
# final_v=output[NT+1, (1+m)*(1+n)+1:2*(m+1)*(n+1)]
# final_p=output[NT+1, 2*(m+1)*(n+1)+1:end]

# u1 = final_u[Int(n/2)*(m+1)+1: Int(n/2)*(m+1)+m+1]
# u2 = final_u[Int(n/2)+1:m+1:end]

# v1 = final_v[Int(n/2)*(m+1)+1: Int(n/2)*(m+1)+m+1]
# v2 = final_v[Int(n/2)+1:m+1:end]
# xx = 0:h:1

# figure();plot(xx, u1);#plot(xx, u12);
# savefig("u_horizontal.png")

# figure();plot(xx, u2);#plot(xx, u22);
# savefig("u_vertical.png")

# figure();plot(xx, v1);#plot(xx, v12);
# savefig("v_horizontal.png")

# figure();plot(xx, v2);#plot(xx, v22);
# savefig("v_vertical.png")

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

# p12 = final_p2[(Int(n/2)-1)*m+1: (Int(n/2)-1)*m+m]
# p22 = final_p2[Int(n/2)*m+1: Int(n/2)*m+m]
# p32 = 0.5 * (p12 .+ p22)

# p42 = final_p2[Int(n/2):m:end]
# p52 = final_p2[Int(n/2)+1:m:end]
# p62 = 0.5 * (p42 .+ p52)

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

# p1 = final_p[(Int(n/2)-1)*m+1: (Int(n/2)-1)*m+m]
# p2 = final_p[Int(n/2)*m+1: Int(n/2)*m+m]
# p3 = 0.5 * (p1 .+ p2)

# p4 = final_p[Int(n/2):m:end]
# p5 = final_p[Int(n/2)+1:m:end]
# p6 = 0.5 * (p4 .+ p5)

# xx = 0.005:h:1

# figure();plot(xx, p3);#plot(xx, p32);
# savefig("p_horizontal.png")

# figure();plot(xx, p6);#plot(xx, p62);
# savefig("p_vertical.png")


"""
Функция сложения
"""
function summa(a::Float64,b::Float64,c::Float64)
    d = round(a+b+c)
end

function summa(a::Int64,b::Int64,c::Int64)
    d = a+b+c
end

function summa(a,b,c)
    return "Invalid input values"
end

"""
#Блабабла
Функция расчета квадратных корней
"""
function calcQuadRoots(a,b,c)
    D = b^2 - 4*a*c
    if D<0
        x1 = (-b+sqrt(Complex(D)))/2/a
        x2 = (-b-sqrt(Complex(D)))/2/a
    else
        x1 = (-b+sqrt(D))/2/a
        x2 = (-b-sqrt(D))/2/a
    end
    return x1,x2
end


# ===============================
# Test minimal (default) options
# ===============================
using Reachability

# linear ODE: x' = Ax
A = [ 0.0509836  0.168159  0.95246   0.33644
      0.42377    0.67972   0.129232  0.126662
      0.518654   0.981313  0.489854  0.588326
      0.38318    0.616014  0.518412  0.778765]

# initial set
X0 = BallInf(ones(4), 0.1)

# default options (computes all variables)
s = solve(InitialValueProblem(LinearContinuousSystem(A), X0), :T=>0.1)

# two variables and custom partition
s = solve(ContinuousSystem(A, X0), :T=>0.1, :partition=>[1:2, 3:4], :vars=>[1,3])

# the default partition is used. uses Interval for set_type_init and set_type_iter
# but Hyperrectangle for set_type_proj
s = solve(ContinuousSystem(A, X0), :T=>0.1, :vars=>[1,3])

# ===============================
# Test projection
# ===============================
s = solve(ContinuousSystem(A, X0), :T=>0.1, :partition=>[1:2, 3:4],
          :vars=>[1,3], :project_reachset=>false);
ps = project(s);

# ===============
# Test check mode
# ===============

A = [ 0.0509836  0.168159  0.95246   0.33644
      0.42377    0.67972   0.129232  0.126662
      0.518654   0.981313  0.489854  0.588326
      0.38318    0.616014  0.518412  0.778765]

# check that x1 + x2 <= 2 doesn't hold
s = solve(ContinuousSystem(A, X0), :T=>0.1, :mode=>"check",
          :partition=>[1:2, 3:4], :vars=>[1,2,3],
          :property=>LinearConstraintProperty([1., 1., 0., 0.], 2.))
@test s.violation == 1

# check that x1 - x2 <= 2 holds
s = solve(ContinuousSystem(A, X0), :T=>0.1, :mode=>"check",
          :partition=>[1:2, 3:4], :vars=>[1,2,3],
          :property=>LinearConstraintProperty([1., -1., 0., 0.], 2.))
@test s.violation == -1

# ===============================
# Test reachability options
# ===============================
s = solve(ContinuousSystem(A, X0), :T=>0.1, :partition=>[1:2, 3:4],
          :vars=>[1,3], :lazy_sih=>true);
s = solve(ContinuousSystem(A, X0), :T=>0.1, :partition=>[1:2, 3:4],
          :vars=>[1,3], :lazy_sih=>false);

s = solve(ContinuousSystem(sparse(A), X0), :T=>0.1, :partition=>[1:2, 3:4],
          :vars=>[1,3], :lazy_expm=>true, :assume_sparse=>false);
s = solve(ContinuousSystem(sparse(A), X0), :T=>0.1, :partition=>[1:2, 3:4],
          :vars=>[1,3], :lazy_expm=>true, :assume_sparse=>true);

s = solve(ContinuousSystem(sparse(A), X0), :T=>0.1, :partition=>[[i] for i in 1:4],
          :set_type=>Interval, :vars=>[1,3], :lazy_expm=>true);

# ===============================
# System with an odd dimension
# ===============================
A = randn(5, 5); X0 = BallInf(ones(5), 0.1)
system = ContinuousSystem(A, X0)
options = Options(Dict(:T=>0.1, :partition=>[1:2, 3:4, [5]], :block_types=>Dict(HPolygon=>[1:2, 3:4], Interval=>[[5]]), :vars=>[1,3]))
s = solve(system, options)

# ===============================
# System with an odd dimension
# ===============================
A = [1 2 1.; 0 0. 1; -2 1 4]
X0 = BallInf(ones(3), 0.1)
system = ContinuousSystem(A, X0)
s = solve(system, :T=>0.1, :vars=>1:3, :partition=>[1:2, 3:3])


@testset "test_conversion" begin

# Easy second order system
sys1 = ss([-1 0;1 1],[1;0],[1 1],0)
G1 = tf([1,0],[1,0,-1])
H1 = zpk([0.0], [-1, 1.0], 1.0)

w = 10.0 .^ (-2:2:50)

@test ss(sys1) == sys1
@test freqresp(ss(G1), w) ≈ freqresp(sys1, w) rtol=1e-15
@test freqresp(ss(H1), w) ≈ freqresp(sys1, w) rtol=1e-15

# Test range in freqresp
@test freqresp(ss(G1), collect(1:2:10)) == freqresp(ss(G1), collect(1:2:10))

@test tf(sys1) ≈ G1 rtol=1e-15
@test tf(G1) == G1
@test tf(H1) == G1

@test zpk(sys1) ≈ H1 rtol=1e-15
@test zpk(G1) == H1
@test zpk(H1) == H1

# With direct term second order system
sys2 = ss([-2.5 0;1 1.5],[1;3],[1 2],2.5)
G2 = tf([2.5, 9.5, 6.125],[1, 1, -3.75])
H2 = zpk([-2.97703296142690, -0.82296703857310], [1.5, -2.5], 2.5)

w = 10.0 .^ (-2:2:50)

@test ss(sys2) == sys2
@test freqresp(ss(G2), w) ≈ freqresp(sys2, w) rtol=1e-15
@test freqresp(ss(H2), w) ≈ freqresp(sys2, w) rtol=1e-15

@test tf(sys2) ≈ G2 rtol=1e-15
@test tf(G2) == G2
@test tf(H2) ≈ G2 rtol=1e-15

@test zpk(sys2) ≈ H2 rtol=1e-15
@test zpk(G2) ≈ H2 rtol=1e-15
@test zpk(H2) == H2

# Non-proper system
G3 = tf([4.0, 1, 5],[2, 3])
H3 = zpk([(-1+im*sqrt(79))/8, (-1-im*sqrt(79))/8], [-3/2], 2)

@test tf(G3) == G3
@test tf(H3) ≈ G3 rtol=1e-2

@test zpk(G3) ≈ H3 rtol=1e-15
@test zpk(H3) == H3

# Test complex 1
A = [1.0 + im 1; 0 -2-3im]
B = [0;2]
C = [1 0]
D = zeros(1,1);
sys = ss(A,B,C,D)
sys2 = convert(TransferFunction, sys)
w = 10.0 .^ (-2:2:50)
@test freqresp(sys, w) ≈ freqresp(sys2, w)
csort = v -> sort(v, lt = (x,y) -> abs(x) < abs(y))
@test csort(pole(zpk(sys2))) ≈ [1+im, -2.0-3im]

# Test complex 2
sys4 = ss([-1.0+im 0;1 1],[1;0],[1 im],im)
tf(sys4)
G4 = tf([1.0im,2,-2],[1,-1im,-1+im])
H4 = zpk([im*(1 + sqrt(1+2im)), im*(1 - sqrt(1+2im))], [-1+im,1], 1.0im)

w = 10.0 .^ (-2:2:50)

@test ss(sys4) == sys4
@test freqresp(ss(G4), w) ≈ freqresp(sys4, w) rtol=1e-15
@test freqresp(ss(H4), w) ≈ freqresp(sys4, w) rtol=1e-15

@test tf(sys4) ≈ G4 rtol=1e-15
@test tf(G4) == G4
@test tf(H4) ≈ G4 rtol=1e-15

@test zpk(sys4) ≈ H4 rtol=1e-15
@test zpk(G4) ≈ H4 rtol=1e-15
@test zpk(H4) == H4

# Test conversion of Int system with direct term
w = [0.5, 1, 1.5, 2.0]
G = tf([2, 3],[5, 1])
@test freqresp(ss(G), w) ≈ freqresp(G, w)

sys2 = ss([-1 0;1 1],[1;0],[1 1],1)
G2 = tf([1,1,-1],[1,0,-1])
sys3 = ss([0 1;1 0],[0;1],[1 1],0)
G3 = tf([1,1],[1,0,-1])

#Test ss2tf
@test tf(sys2) ≈ G2 atol=1e-15
@test tf(sys3) ≈ G3 atol=1e-15

#Test TransferFunction -> StateSpace
#ss is not unique so go back and forth
@test G2 ≈ tf(ss(G2)) atol=1e-15
@test G3 ≈ tf(ss(G3)) atol=1e-15

### test mimo
@test hinfnorm(zpk([sys1 sys2])-zpk([ss(sys1) ss(sys2)]))[1] < 1e-12

## Test some BigFloat
SSBigFloat = StateSpace{BigFloat,Array{BigFloat,2}}
ZpkBigFloat = TransferFunction{ControlSystems.SisoZpk{BigFloat,Complex{BigFloat}}}
RationalBigFloat = TransferFunction{ControlSystems.SisoRational{BigFloat}}

s = tf("s")
f = zpk(1.0*(2s+3)/((5s+7)*(11s+13)))
fb = BigFloat(1.0)*f

# Reasonable defaults
# @test fb isa RationalBigFloat # Do we want this? We get TransferFunction{ControlSystems.SisoZpk{BigFloat,Complex{Float64}}}
@test ss(fb) isa SSBigFloat
@test tf(fb) isa RationalBigFloat
# Cant compute poles:
# @test zpk(tf(fb)) isa SSBigFloat

@test convert(SSBigFloat, fb) isa SSBigFloat
@test convert(ZpkBigFloat, fb) isa ZpkBigFloat
@test convert(RationalBigFloat, fb) isa RationalBigFloat

## Test some rather trivial conversions of numeric types

b = 1.5
D22 = [1.0 2.0; 3.0 4.0]

@test convert(StateSpace{Float64,Matrix{Float64}}, D22) == ss(D22)
@test convert(StateSpace{Float64,Matrix{Float64}}, b) == ss(b)

@test convert(TransferFunction{ControlSystems.SisoRational{Float64}}, D22) == tf(D22)
@test convert(TransferFunction{ControlSystems.SisoRational{Float64}}, b) == tf(b)

@test convert(TransferFunction{ControlSystems.SisoZpk{Float64,ComplexF64}}, D22) == zpk(D22)
@test convert(TransferFunction{ControlSystems.SisoZpk{Float64,ComplexF64}}, b) == zpk(b)

# Error, not proper
@test_throws ErrorException ss(tf([1,0],[1]))

end


### A Pluto.jl notebook ###
# v0.12.20

using Markdown
using InteractiveUtils

# ╔═╡ 57aa3170-6e51-11eb-2670-9b77065bc4da
begin
    using PlutoUI
    using Plots
    using Printf
	using SpecialFunctions
end

# ╔═╡ ad3d9470-6e82-11eb-3072-43decfc60ef9
md"_MAC Method for Lid-driven Cavity Problem, version 1_"

# ╔═╡ 484fcd20-6f23-11eb-0dc8-5d49530983a1
md"""
## **Lid-driven Cavity Problem with MAC Method**
#### **Description:**
This notebook uses MAC method to solve incompressible 2D Navier-Stokes equations for the solution of the lid-driven cavity problem.

		MAC(Re,h)

If the arguments are missing, it will use the default values.

		MAC(Re=0.002,h=0.002)

#### **Lid-driven Cavity Problem**
In this problem, there is a moving lid over a cavity. The lid is moving with a constant speed. The governing equations for the system can be written in the conservative form:

$$\frac{∂u}{∂x}+\frac{∂u}{∂y}=0$$
$$\frac{∂u}{∂t}+\frac{∂}{∂x}(u^2+\frac{p}{ρ})+\frac{∂(uv)}{∂y}=ν(\frac{∂^2u}{∂x^2}+\frac{∂^2u}{∂y^2})$$
$$\frac{∂v}{∂t}+\frac{∂(uv)}{∂x}(v^2+\frac{p}{ρ})+\frac{∂(uv)}{∂y}=ν(\frac{∂^2v}{∂x^2}+\frac{∂^2v}{∂y^2})$$

These equations can be nondimensionalized with the following values:

$$Re=\frac{ρ_\infty u_\infty L}{μ_\infty}$$
$$t^*=\frac{t u_\infty}{L}$$
$$x^*=\frac{x}{L}$$
$$y^*=\frac{y}{L}$$
$$u^*=\frac{u}{u_\infty}$$
$$v^*=\frac{v}{u_\infty}$$
$$p^*=\frac{p}{ρ_\infty u^2_\infty}$$

After the nondimensionalization * is ignored for simplicity:

$$\frac{∂u}{∂x}+\frac{∂u}{∂y}=0$$
$$\frac{∂u}{∂t}+\frac{∂}{∂x}(u^2+p)+\frac{∂(uv)}{∂y}=\frac{1}{Re}(\frac{∂^2u}{∂x^2}+\frac{∂^2u}{∂y^2})$$
$$\frac{∂v}{∂t}+\frac{∂(uv)}{∂x}(v^2+p)+\frac{∂(uv)}{∂y}=\frac{1}{Re}(\frac{∂^2v}{∂x^2}+\frac{∂^2v}{∂y^2})$$

Since continuity term does not include a time dependent term, there is not communication mechanism between continuity and momentum equations. To overcome this problem, possion equation will be solved. In this notebook, calculations are done in three steps:

1) Calculation of u*.
2) Calculation of p.
3) Calculation of corrected u.

This approach is called pressure correction. The discritization for the system is:

$$I   = \frac{(u_{i+1,j}+u_{i,j})^2-(u_{i,j}+u_{i-1,j})^2}{4Δx}$$
$$II  = \frac{(u_{i,j+1}+u_{i,j})(v_{i+1,j}+v_{i,j})}{4Δy}-\frac{(u_{i,j}+u_{i,j-1})(v_{i+1,j-1}+v_{i,j-1})}{4Δy}$$
$$III = \frac{1}{Re}(\frac{u_{i+1,j}-2u_{i,j}+u_{i-1,j}}{Δx^2}+\frac{u_{i,j+1}-2u_{i,j}+u_{i,j-1}}{Δy^2})$$
$$\frac{u_{i,j}^*-u_{i,j}^n}{Δt}=-I-II+III$$

For the pressure correction:

$$I=\frac{p_{i+1,j}-2p_{i,j}+p_{i-1,j}}{Δx^2}+\frac{p_{i,j+1}-2p_{i,j}+p_{i,j-1}}{Δy^2}$$
$$II=\frac{\frac{u_{i,j}^*-u_{i-1,j}^*}{Δx}+\frac{v_{i,j}^*-v_{i,j-1}^*}{Δy}}{Δt}$$
$$p_{i,j}=p_{i,j}+ω\frac{Δx^2+Δy^2}{2}(I-II)$$

In this notebook, ω is taken as 1 and this correction method is applied 10 times. Once the pressure is calculated, correct velocity can be calculated for the next time step as:

$$\frac{u^{n+1}_{i,j}-u^*_{i,j}}{Δt}=\frac{p_{i+1,j}-p_{i,j}}{Δx}$$

Same procedure can be applied for the v velocity. Boundary conditions for the system is:

$$u(0,:)=0$$
$$u(N,:)=0$$
$$u(:,0)=0$$
$$u(:,N)=1$$
$$v(0,:)=0$$
$$v(N,:)=0$$
$$v(:,0)=0$$
$$v(:,N)=0$$

For the pressure Nuemann boundary condition is applied as:

$$\frac{∂p}{∂η}=0$$

where η is the wall normal direction.

Details of schemes:
Computational Fluid Dynamics, Volume I, Hoffmann - Chiang

Feel free to ask questions!

*Furkan Oz*
foz@okstate.edu
"""

# ╔═╡ 6bc808d0-6e51-11eb-0794-b9585413d41b

"""
This notebook uses MAC method to solve incompressible 2D Navier-Stokes equations for the solution of the lid-driven cavity problem.

		MAC(Re,h)

If the arguments are missing, it will use the default values.

		MAC(Re=0.002,h=0.002)

Furkan Oz,
foz@okstate.edu,

"""
function MAC(Re=100, h=0.02)
	# Initialiazing for domain setup
	xdom = [0 1.0] # Domain limits for x
	ydom = [0 1.0] # Domain limits for y
	Δx = h
	Δy = h
	Δx²= Δx^2
	Δy²= Δy^2
	nex = floor(Int32,(xdom[2]-xdom[1])/h) #Number of elements in x-direction
	ney = floor(Int32,(xdom[2]-xdom[1])/h) #Number of elements in y-direction
	nnx = nex + 1 #Number of nodes in x-direction
	nny = ney + 1 #Number of nodes in y-direction

	# Initialiazing for solver setup

	CFL = 1 # CFL condition (less than 1 preferred)
	Δtc = CFL*min(Δx,Δy) # Time-step based on Convective Term
	Δtv = 0.25*Re*min(Δx²,Δy²) # Time-step based on Viscous Term
	σ = 0.5
	a =-2.0/(Δx²) - 2.0/(Δy²)
	Δt =  σ*min(Δtc,Δtv) # Time-step based on lowest condition
	ϵ = 1e-5 # Convergence limit

	ttmax = 30000 # Maximum iteration limit
	iimax = 10 # Pressure correction iteration limit

	u  = zeros(nnx,nny+1)
	v  = zeros(nnx+1,nny)
	uˢ = zeros(nnx,nny+1)
	vˢ = zeros(nnx+1,nny)
	u₀ = zeros(nnx,nny+1)
	v₀ = zeros(nnx+1,nny)

	p  = zeros(nex+2,ney+2)
	p₀ = zeros(nex+2,ney+2)

	rᵤ = 0
	rᵥ = 0
	rₚ = 0
	ttend = 0

	rᵤₜ = zeros(ttmax)
	rᵥₜ = zeros(ttmax)
	rₚₜ = zeros(ttmax)

	#Solver starts

	for tt=1:ttmax
		for j=1:nny+1
			for i=1:nnx
				u₀[i,j] = u[i,j]
			end
		end
		for j=1:nny
			for i=1:nnx+1
				v₀[i,j] = v[i,j]
			end
		end
		for j=1:ney+2
			for i=1:nex+2
				p₀[i,j] = p[i,j]
			end
		end

		# Boundary conditions for velocity
		for i=1:nnx
			u[i,1] =-u[i,2]
			u[i,nny+1] = 2.0-u[i,nny]
		end
		for i=1:nny
			v[1,i] =-v[2,i]
			v[nnx+1,i] =-v[nnx,i]
		end

		# uˢ calculation
		for j=2:nny
			for i=2:nnx-1

  				c₃ = (1.0/Re)*((u[i+1,j]-2.0*u[i,j]+u[i-1,j])/(Δx²)  				                      + (u[i,j+1]-2.0*u[i,j]+u[i,j-1])/(Δy²) )

  				c₁ = (0.25/Δx)*((u[i+1,j]+u[i,j])^2 - (u[i,j]+u[i-1,j])^2)

  				c₂ = (0.25/Δy)*((u[i,j+1]+u[i,j])*(v[i+1,j]+v[i,j])                                      - (u[i,j]+u[i,j-1])*(v[i+1,j-1]+v[i,j-1]))

  				uˢ[i,j] = u[i,j] + Δt*(-c₁-c₂+c₃)

			end
		end

		# vˢ calculation
		for j=2:nny-1
			for i=2:nnx

  				c₃ = (1.0/Re)*((v[i+1,j]-2.0*v[i,j]+v[i-1,j])/(Δx²)  				                      + (v[i,j+1]-2.0*v[i,j]+v[i,j-1])/(Δy²) )

  				c₂ = (0.25/Δy)*((v[i,j+1]+v[i,j])^2 - (v[i,j]+v[i,j-1])^2)

  				c₁ = (0.25/Δx)*((u[i,j+1]+u[i,j])*(v[i+1,j]+v[i,j])                                      - (u[i-1,j+1]+u[i-1,j])*(v[i,j]+v[i-1,j]))

  				vˢ[i,j] = v[i,j] + Δt*(-c₁-c₂+c₃)

			end
		end

		# Boundary conditions for velocity correction
		for i=1:nnx
			uˢ[i,1] =-uˢ[i,2]
			uˢ[i,nny+1] = 2.0-uˢ[i,nny]
		end
		for i=1:nny
			vˢ[1,i] =-vˢ[2,i]
			vˢ[nnx+1,i] =-vˢ[nnx,i]
		end

		# Boundary conditions for pressure
		for i=1:nex+2
			p[i,1] = p[i,2]
			p[i,ney+2] = p[i,ney+1]
		end
		for i=1:ney+2
			p[1,i] = p[2,i]
			p[nex+2,i] = p[nex+1,i]
		end

		# Poisson Solver
		ω = 1.0
		for ii = 1:iimax
			for j=2:ney+1
				for i=2:nex+1

  					f = ((uˢ[i,j]-uˢ[i-1,j])/Δx + (vˢ[i,j]-vˢ[i,j-1])/Δy)/Δt

  					r = (f - (p[i+1,j]-2.0*p[i,j]+p[i-1,j])/(Δx²)                                           -      (p[i,j+1]-2.0*p[i,j]+p[i,j-1])/(Δy²) )

  					p[i,j] = p[i,j] + ω*r/a

				end
			end
		end

		# Correct velocity calculation
		for j=2:nny
			for i=2:nnx-1
  				u[i,j] = uˢ[i,j] - Δt*(p[i+1,j]-p[i,j])/Δx
			end
		end

		for j=2:nny-1
			for i=2:nnx
  				v[i,j] = vˢ[i,j] - Δt*(p[i,j+1]-p[i,j])/Δy
			end
		end

		#check residual
		rᵤ = 0.0
		rᵥ = 0.0
		rₚ = 0.0
		for j=2:nny
			for i=2:nnx
  				rᵤ = rᵤ + (u[i,j] - u₀[i,j])^2
  				rᵥ = rᵥ + (v[i,j] - v₀[i,j])^2
  				rₚ = rₚ + (p[i,j] - p₀[i,j])^2
			end
		end
		rᵤ = sqrt(rᵤ/(nex*ney))
		rᵥ = sqrt(rᵥ/(nex*ney))
		rₚ = sqrt(rₚ/(nex*ney))

		rᵤₜ[tt] = rᵤ
		rᵥₜ[tt] = rᵥ
		rₚₜ[tt] = rₚ
		ttend = tt
		@show tt rᵤ rᵥ rₚ ttend rᵤₜ[tt]
		if (rᵤ<=ϵ && rᵥ<=ϵ && rₚ<=ϵ)
			break
		end

	end
	ttend = 1:ttend

	return u,v,p,nnx,ttend,rᵤₜ,rᵥₜ,rₚₜ
end

# ╔═╡ 9be09b70-6f2a-11eb-0f64-018b8c98e73f
begin
	u100c,v100c,p100c = MAC(100)
	contourf(u100c',levels=50,c= :redsblues, clims=(-1, 1),linewidth = 0,title = "Re=100 - u Contour")
end

# ╔═╡ 6c9f9cc0-6e54-11eb-380f-9f8f2d11343c
begin
	Re100data = [1	1	0.9766	0.84123	0.9688	0.78871	0.9609	0.73722	0.9531	0.68717	0.8516	0.23151	0.7344	0.00332	0.6172	-0.13641	0.5000 -0.20581	0.4531 -0.21090	0.2813	-0.15662	0.1719	-0.1015	0.1016	-0.06434	0.0703	-0.04775	0.0625	-0.04192	0.0547	-0.03717	0	0]
	y100data = [Re100data[i] for i=1:2:34]
	u100data = [Re100data[i] for i=2:2:34]

	u100_0_02,v100_0_02,p100_0_02,nnx100_0_02 = MAC(100,0.02)
	y100_0_02 = -0.01:0.02:1+0.01
	nnx100_0_02m = floor(Int32,(nnx100_0_02-1)/2)

	u100_0_01,v100_0_01,p100_0_01,nnx100_0_01 = MAC(100,0.01)
	y100_0_01 = -0.005:0.01:1+0.005
	nnx100_0_01m = floor(Int32,(nnx100_0_01-1)/2)

	u100_0_005,v100_0_005,p100_0_005,nnx100_0_005 = MAC(100,0.005)
	y100_0_005 = -0.0025:0.005:1+0.0025
	nnx100_0_005m = floor(Int32,(nnx100_0_005-1)/2)

	plot([u100_0_02[nnx100_0_02m,:], u100_0_01[nnx100_0_01m,:], u100_0_005[nnx100_0_005m,:], u100data],[y100_0_02,y100_0_01,y100_0_005,y100data],
				title = "Re=100 - Center u Velocity",
				label = ["MAC h=0.02" "MAC h=0.01" "MAC h=0.005" "Ghia"],
				legend = :bottomright,
				#ylims = (0,0.04),
				xlabel = "U",
				ylabel = "y",
				linewidth = 2,
				linecolor = [:black :black :black :transparent],
				markershape = :auto,
				markercolor = [:red :red :red :green],
				)
end

# ╔═╡ 07e46510-6f1f-11eb-21a0-f50e1e1e67b8
begin

	 x,x,x,x,tt100_0_02,ru100_0_02,rv100_0_02,rp100_0_02 = MAC(100,0.02)

	plot([tt100_0_02],[log.(ru100_0_02[1:tt100_0_02[end]]),log.(rv100_0_02[1:tt100_0_02[end]]),log.(rp100_0_02[1:tt100_0_02[end]])],
				title = "Re=100 - h=0.02 - Residual",
				label = ["u residual" "v residual" "p residual"],
				legend = :topright,
				xlims = (0,1500),
				ylims = (-15,-2.5),
				xlabel = "Iteration",
				ylabel = "log(Residual)",
				linewidth = 2,
				linecolor = [:auto],
				#markershape = :auto,
				#markercolor = [:red],
				)
end

# ╔═╡ a71f5160-6f26-11eb-1ffb-15467ef1e901
begin

	q,q,q,q,tt100_0_01,ru100_0_01,rv100_0_01,rp100_0_01 = MAC(100,0.01)

	plot([tt100_0_01],[log.(ru100_0_01[1:tt100_0_01[end]]),log.(rv100_0_01[1:tt100_0_01[end]]),log.(rp100_0_01[1:tt100_0_01[end]])],
				title = "Re=100 - h=0.01 - Residual",
				label = ["u residual" "v residual" "p residual"],
				legend = :topright,
				#xlims = (0,500),
				ylims = (-15,-2.5),
				xlabel = "Iteration",
				ylabel = "log(Residual)",
				linewidth = 2,
				linecolor = [:auto],
				#markershape = :auto,
				#markercolor = [:red],
				)
end

# ╔═╡ 13382bd0-6f2a-11eb-392a-6fc93389bda6
begin

	w,w,w,w,tt100_0_005,ru100_0_005,rv100_0_005,rp100_0_005 = MAC(100,0.005)

	plot([tt100_0_005],[log.(ru100_0_005[1:tt100_0_005[end]]),log.(rv100_0_005[1:tt100_0_005[end]]),log.(rp100_0_005[1:tt100_0_005[end]])],
				title = "Re=100 - h=0.005 - Residual",
				label = ["u residual" "v residual" "p residual"],
				legend = :topright,
				#xlims = (0,500),
				ylims = (-15,-2.5),
				xlabel = "Iteration",
				ylabel = "log(Residual)",
				linewidth = 2,
				linecolor = [:auto],
				#markershape = :auto,
				#markercolor = [:red],
				)
end

# ╔═╡ 4c463b00-6f2b-11eb-2840-7384b467a27b
begin
	u400c,v400c,p400c = MAC(400)
	contourf(u400c',levels=50,c= :redsblues, clims=(-1, 1),linewidth = 0,title = "Re=400 - u Contour")
end

# ╔═╡ bfad9c20-6e7f-11eb-2fc9-49f1df970d96
begin
	Re400data = [1	1	0.9766	0.75837	0.9688	0.68439	0.9609	0.61756	0.9531	0.55892	0.8516	0.29093	0.7344	0.16256	0.6172	0.02135	0.5000 -0.11477	0.4531 -0.17119	0.2813	-0.32726	0.1719	-0.24299	0.1016	-0.14612	0.0703	-0.10338	0.0625	-0.09266	0.0547	-0.08186	0	0
]
	y400data = [Re400data[i] for i=1:2:34]
	u400data = [Re400data[i] for i=2:2:34]

	u400_0_02,v400_0_02,p400_0_02,nnx400_0_02 = MAC(400,0.02)
	y400_0_02 = -0.01:0.02:1+0.01
	nnx400_0_02m = floor(Int32,(nnx400_0_02-1)/2)

	u400_0_01,v400_0_01,p400_0_01,nnx400_0_01 = MAC(400,0.01)
	y400_0_01 = -0.005:0.01:1+0.005
	nnx400_0_01m = floor(Int32,(nnx400_0_01-1)/2)

	u400_0_005,v400_0_005,p400_0_005,nnx400_0_005 = MAC(400,0.005)
	y400_0_005 = -0.0025:0.005:1+0.0025
	nnx400_0_005m = floor(Int32,(nnx400_0_005-1)/2)

	plot([u400_0_02[nnx400_0_02m,:], u400_0_01[nnx400_0_01m,:], u400_0_005[nnx400_0_005m,:], u400data],[y400_0_02,y400_0_01,y400_0_005,y400data],
				title = "Re=400 - Center Velocity",
				label = ["MAC h=0.02" "MAC h=0.01" "MAC h=0.005" "Ghia"],
				legend = :bottomright,
				#ylims = (0,0.04),
				xlabel = "U",
				ylabel = "y",
				linewidth = 2,
				linecolor = [:black :black :black :transparent],
				markershape = :auto,
				markercolor = [:red :red :red :green],
				)
end

# ╔═╡ 26043960-6f2b-11eb-2fdd-33f7de7fc9e3
begin

	 t,t,t,t,tt400_0_02,ru400_0_02,rv400_0_02,rp400_0_02 = MAC(400,0.02)

	plot([tt400_0_02],[log.(ru400_0_02[1:tt400_0_02[end]]),log.(rv400_0_02[1:tt400_0_02[end]]),log.(rp400_0_02[1:tt400_0_02[end]])],
				title = "Re=400 - h=0.02 - Residual",
				label = ["u residual" "v residual" "p residual"],
				legend = :topright,
				xlims = (0,1500),
				ylims = (-15,-2.5),
				xlabel = "Iteration",
				ylabel = "log(Residual)",
				linewidth = 2,
				linecolor = [:auto],
				#markershape = :auto,
				#markercolor = [:red],
				)
end

# ╔═╡ 03cf5b90-6f2b-11eb-0dd8-7f2061f9d9d7
begin

	r,r,r,r,tt400_0_01,ru400_0_01,rv400_0_01,rp400_0_01 = MAC(400,0.01)

	plot([tt400_0_01],[log.(ru400_0_01[1:tt400_0_01[end]]),log.(rv400_0_01[1:tt400_0_01[end]]),log.(rp400_0_01[1:tt400_0_01[end]])],
				title = "Re=400 - h=0.01 - Residual",
				label = ["u residual" "v residual" "p residual"],
				legend = :topright,
				#xlims = (0,500),
				ylims = (-15,-2.5),
				xlabel = "Iteration",
				ylabel = "log(Residual)",
				linewidth = 2,
				linecolor = [:auto],
				#markershape = :auto,
				#markercolor = [:red],
				)
end

# ╔═╡ e680ea90-6f2a-11eb-1ca2-89a396add38a
begin

	e,e,e,e,tt400_0_005,ru400_0_005,rv400_0_005,rp400_0_005 = MAC(400,0.005)

	plot([tt400_0_005],[log.(ru400_0_005[1:tt400_0_005[end]]),log.(rv400_0_005[1:tt400_0_005[end]]),log.(rp400_0_005[1:tt400_0_005[end]])],
				title = "Re=400 - h=0.005 - Residual",
				label = ["u residual" "v residual" "p residual"],
				legend = :topright,
				#xlims = (0,500),
				ylims = (-15,-2.5),
				xlabel = "Iteration",
				ylabel = "log(Residual)",
				linewidth = 2,
				linecolor = [:auto],
				#markershape = :auto,
				#markercolor = [:red],
				)
end

# ╔═╡ 57ed04c2-6f2b-11eb-20c8-1725c926dd2a
begin
	u1000c,v1000c,p1000c = MAC(1000)
	contourf(u1000c',levels=50,c= :redsblues, clims=(-1, 1),linewidth = 0,title = "Re=1000 - u Contour")
end

# ╔═╡ c8ece650-6e80-11eb-2882-7fed169d61ed
begin
	Re1000data = [1	1	0.9766	0.65928	0.9688	0.57492	0.9609	0.51117	0.9531	0.46604	0.8516	0.33304	0.7344	0.18719	0.6172	0.05702	0.5	-0.0608	0.4531	-0.10648	0.2813	-0.27805	0.1719	-0.38289	0.1016	-0.2973	0.0703	-0.2222	0.0625	-0.20196	0.0547	-0.18109	0	0]
	y1000data = [Re1000data[i] for i=1:2:34]
	u1000data = [Re1000data[i] for i=2:2:34]

	u1000_0_02,v1000_0_02,p1000_0_02,nnx1000_0_02 = MAC(1000,0.02)
	y1000_0_02 = -0.01:0.02:1+0.01
	nnx1000_0_02m = floor(Int32,(nnx1000_0_02-1)/2)

	u1000_0_01,v1000_0_01,p1000_0_01,nnx1000_0_01 = MAC(1000,0.01)
	y1000_0_01 = -0.005:0.01:1+0.005
	nnx1000_0_01m = floor(Int32,(nnx1000_0_01-1)/2)

	u1000_0_005,v1000_0_005,p1000_0_005,nnx1000_0_005 = MAC(1000,0.005)
	y1000_0_005 = -0.0025:0.005:1+0.0025
	nnx1000_0_005m = floor(Int32,(nnx1000_0_005-1)/2)

	plot([u1000_0_02[nnx1000_0_02m,:], u1000_0_01[nnx1000_0_01m,:], u1000_0_005[nnx1000_0_005m,:], u1000data],[y1000_0_02,y1000_0_01,y1000_0_005,y1000data],
				title = "Re=1000 - Center Velocity",
				label = ["MAC h=0.02" "MAC h=0.01" "MAC h=0.005" "Ghia"],
				legend = :bottomright,
				#ylims = (0,0.04),
				xlabel = "U",
				ylabel = "y",
				linewidth = 2,
				linecolor = [:black :black :black :transparent],
				markershape = :auto,
				markercolor = [:red :red :red :green],
				)
end

# ╔═╡ 94b982e0-6f29-11eb-2f68-839a418ad1f8
begin

	 c,c,c,c,tt1000_0_02,ru1000_0_02,rv1000_0_02,rp1000_0_02 = MAC(1000,0.02)

	plot([tt1000_0_02],[log.(ru1000_0_02[1:tt1000_0_02[end]]),log.(rv1000_0_02[1:tt1000_0_02[end]]),log.(rp1000_0_02[1:tt1000_0_02[end]])],
				title = "Re=1000 - h=0.02 - Residual",
				label = ["u residual" "v residual" "p residual"],
				legend = :topright,
				#xlims = (0,1500),
				ylims = (-15,-2.5),
				xlabel = "Iteration",
				ylabel = "log(Residual)",
				linewidth = 2,
				linecolor = [:auto],
				#markershape = :auto,
				#markercolor = [:red],
				)

end

# ╔═╡ b5637060-6f28-11eb-0850-5dff9a108c4b
begin

	 b,b,b,b,tt1000_0_01,ru1000_0_01,rv1000_0_01,rp1000_0_01 = MAC(1000,0.01)

	plot([tt1000_0_01],[log.(ru1000_0_01[1:tt1000_0_01[end]]),log.(rv1000_0_01[1:tt1000_0_01[end]]),log.(rp1000_0_01[1:tt1000_0_01[end]])],
				title = "Re=1000 - h=0.01 - Residual",
				label = ["u residual" "v residual" "p residual"],
				legend = :topright,
				#xlims = (0,1500),
				ylims = (-15,-2.5),
				xlabel = "Iteration",
				ylabel = "log(Residual)",
				linewidth = 2,
				linecolor = [:auto],
				#markershape = :auto,
				#markercolor = [:red],
				)

end

# ╔═╡ 2e0fecf0-6f29-11eb-1d52-71a1c891ed5a
begin

	 a,a,a,a,tt1000_0_005,ru1000_0_005,rv1000_0_005,rp1000_0_005 = MAC(1000,0.005)

	plot([tt1000_0_005],[log.(ru1000_0_005[1:tt1000_0_005[end]]),log.(rv1000_0_005[1:tt1000_0_005[end]]),log.(rp1000_0_005[1:tt1000_0_005[end]])],
				title = "Re=1000 - h=0.005 - Residual",
				label = ["u residual" "v residual" "p residual"],
				legend = :topright,
				#xlims = (0,1500),
				ylims = (-15,-2.5),
				xlabel = "Iteration",
				ylabel = "log(Residual)",
				linewidth = 2,
				linecolor = [:auto],
				#markershape = :auto,
				#markercolor = [:red],
				)

end

# ╔═╡ Cell order:
# ╟─57aa3170-6e51-11eb-2670-9b77065bc4da
# ╟─ad3d9470-6e82-11eb-3072-43decfc60ef9
# ╟─484fcd20-6f23-11eb-0dc8-5d49530983a1
# ╟─6bc808d0-6e51-11eb-0794-b9585413d41b
# ╟─9be09b70-6f2a-11eb-0f64-018b8c98e73f
# ╟─6c9f9cc0-6e54-11eb-380f-9f8f2d11343c
# ╟─07e46510-6f1f-11eb-21a0-f50e1e1e67b8
# ╟─a71f5160-6f26-11eb-1ffb-15467ef1e901
# ╟─13382bd0-6f2a-11eb-392a-6fc93389bda6
# ╟─4c463b00-6f2b-11eb-2840-7384b467a27b
# ╟─bfad9c20-6e7f-11eb-2fc9-49f1df970d96
# ╟─26043960-6f2b-11eb-2fdd-33f7de7fc9e3
# ╟─03cf5b90-6f2b-11eb-0dd8-7f2061f9d9d7
# ╟─e680ea90-6f2a-11eb-1ca2-89a396add38a
# ╟─57ed04c2-6f2b-11eb-20c8-1725c926dd2a
# ╟─c8ece650-6e80-11eb-2882-7fed169d61ed
# ╟─94b982e0-6f29-11eb-2f68-839a418ad1f8
# ╟─b5637060-6f28-11eb-0850-5dff9a108c4b
# ╟─2e0fecf0-6f29-11eb-1d52-71a1c891ed5a


# Functions and operators on models
"""
    _model_dicts(model::Model)

    input: a Model object
outputs:
     - SDict: a dictionary with keys `"+","-","0","bullet"`
    and corresponding values `findall(model.C .> 0)`, `findall(model.C .< 0)`,
    `findall(model.C .== 0)`, `findall(model.C .!= 0)`, respectively.
     - TDict: a dictionary of submatrices of `T` with keys
    `"ℓm"` with ``ℓ,m∈{+,-,0,bullet}`` and corresponding values
    `model.T[S[ℓ],S[m]]`.
"""
function _model_dicts(model::Model)
    nPhases = n_phases(model)
    SDict = Dict{String,Array}("S" => 1:nPhases)
    SDict["+"] = findall(rates(model) .> 0)
    SDict["-"] = findall(rates(model) .< 0)
    SDict["0"] = findall(rates(model) .== 0)
    SDict["bullet"] = findall(rates(model) .!= 0)

    TDict = Dict{String,Array}("T" => model.T)
    for ℓ in ["+" "-" "0" "bullet"], m in ["+" "-" "0" "bullet"]
        TDict[ℓ*m] = model.T[SDict[ℓ], SDict[m]]
    end

    return SDict, TDict
end

"""
Construct and evaluate ``Ψ(s)`` for a triditional SFM.

Uses newtons method to solve the Ricatti equation
``D⁺⁻(s) + Ψ(s)D⁻⁺(s)Ψ(s) + Ψ(s)D⁻⁻(s) + D⁺⁺(s)Ψ(s) = 0.``

    psi_fun_x( model::Model; s = 0, MaxIters = 1000, err = 1e-8)

# Arguments
- `model`: a Model object
- `s::Real`: a value to evaluate the LST at
- `MaxIters::Int`: the maximum number of iterations of newtons method
- `err::Float64`: an error tolerance for terminating newtons method. Terminates
    when `max(Ψ_{n} - Ψ{n-1}) .< eps`.

# Output
- `Ψ(s)::Array{Float64,2}` the matrix ``Ψ``
"""
function psi_fun_x( model::Model; s = 0, MaxIters = 1000, err = 1e-8)
    SDict, TDict = _model_dicts(model)

    T00inv = inv(TDict["00"] - s * LinearAlgebra.I)
    # construct the generator Q(s)
    Q =
        (1 ./ abs.(rates(model)[SDict["bullet"]])) .* (
            TDict["bulletbullet"] - s * LinearAlgebra.I -
            TDict["bullet0"] * T00inv * TDict["0bullet"]
        )

    model_without_zero_phases = BoundedFluidQueue(Q,model.S[SDict["bullet"]],model.b)

    ~, QDict = _model_dicts(model_without_zero_phases)

    Ψ = zeros(Float64, length(SDict["+"]), length(SDict["-"]))
    A = QDict["++"]
    B = QDict["--"]
    D = QDict["+-"]
    # use netwons method to solve the Ricatti equation
    for n in 1:MaxIters
        Ψ = LinearAlgebra.sylvester(A,B,D)
        if maximum(abs.(sum(Ψ,dims=2).-1)) < err
            break
        end
        A = QDict["++"] + Ψ * QDict["-+"]
        B = QDict["--"] + QDict["-+"] * Ψ
        D = QDict["+-"] - Ψ * QDict["-+"] * Ψ
    end

    return Ψ
end

"""
Construct the vector ``ξ`` containing the distribution of the phase at the time
when ``X(t)`` first hits `0`.

    xi_x( model::Model, Ψ::Array)

# Arguments
- `model`: a Model object
- `Ψ`: an array as output from `psi_fun_x`

# Output
- the vector `ξ`
"""
function xi_x( model::Model, Ψ::Array)
    # the system to solve is [ξ 0](-[B₋₋ B₋₀; B₀₋ B₀₀])⁻¹[B₋₊; B₀₊]Ψ = ξ
    # writing this out and using block inversion (as described on wikipedia)
    # we can solve this in the following way
    SDict, TDict = _model_dicts(model)

    T00inv = inv(TDict["00"])
    invT₋₋ =
        inv(TDict["--"] - TDict["-0"] * T00inv * TDict["0-"])
    invT₋₀ = -invT₋₋ * TDict["-0"] * T00inv

    A =
        -(
            invT₋₋ * TDict["-+"] * Ψ + invT₋₀ * TDict["0+"] * Ψ +
            LinearAlgebra.I
        )
    b = zeros(1, size(TDict["--"], 1))
    A[:, 1] .= 1.0 # normalisation conditions
    b[1] = 1.0 # normalisation conditions

    ξ = b / A

    return ξ
end

"""
Construct the stationary distribution of the SFM

    stationary_distribution_x( model::Model, Ψ::Array, ξ::Array)

# Arguments
- `model`: a Model object
- `Ψ`: an array as output from `psi_fun_x`
- `ξ`: an array as returned from `xi_x`

# Output
- `pₓ::Array{Float64,2}`: the point masses of the SFM
- `πₓ(x)` a function with two methods
    - `πₓ(x::Real)`: for scalar inputs, returns the stationary density evaluated
        at `x` in all phases.
    - `πₓ(x::Array)`: for array inputs, returns an array with the same shape
        as is output by Coeff2Dist.
- `K::Array{Float64,2}`: the matrix in the exponential of the density.
"""
function stationary_distribution_x( model::Model, Ψ::Array, ξ::Array)
    # using the same block inversion trick as in xi_x
    SDict, TDict = _model_dicts(model)

    T00inv = inv(TDict["00"])
    invT₋₋ =
        inv(TDict["--"] - TDict["-0"] * T00inv * TDict["0-"])
    invT₋₀ = -invT₋₋ * TDict["-0"] * T00inv

    Q =
        (1 ./ abs.(rates(model)[SDict["bullet"]])) .* (
            TDict["bulletbullet"] -
            TDict["bullet0"] * T00inv * TDict["0bullet"]
        )

    model_without_zero_phases = BoundedFluidQueue(Q,model.S[SDict["bullet"]],model.b)

    ~, QDict = _model_dicts(model_without_zero_phases)

    K = QDict["++"] + Ψ * QDict["-+"]

    A = -[invT₋₋ invT₋₀]

    # unnormalised values
    αpₓ = ξ * A

    απₓ = αpₓ *
        [TDict["-+"]; TDict["0+"]] *
        -inv(K) *
        [LinearAlgebra.I(length(SDict["+"])) Ψ] *
        LinearAlgebra.diagm(1 ./ abs.(rates(model)[SDict["bullet"]]))

    απₓ0 = -απₓ * [TDict["+0"];TDict["-0"]] * T00inv

    # normalising constant
    α = sum(αpₓ) + sum(απₓ) + sum(απₓ0)

    # normalised values
    # point masses
    pₓ = αpₓ/α
    # density method for scalar x-values
    idx = [findall(rates(model).>0);
        findall(rates(model).<0);
        findall(rates(model).==0)]
    function πₓ(x::Real)
        out = zeros(n_phases(model))
        out[idx] = (pₓ *
        [TDict["-+"]; TDict["0+"]] *
        exp(K*x) *
        [LinearAlgebra.I(length(SDict["+"])) Ψ] *
        LinearAlgebra.diagm(1 ./ abs.(rates(model)[SDict["bullet"]])) *
        [LinearAlgebra.I(sum(rates(model) .!= 0)) [TDict["+0"];TDict["-0"]] * -T00inv])
        return out
    end
    # density method for arrays so that πₓ returns an array with the same shape
    # as is output by Coeff2Dist
    function πₓ(x::Array)
        temp = πₓ.(x)
        Evalπₓ = zeros(Float64, size(x,1), size(x,2), n_phases(model))
        for cell in 1:size(x,2)
            for basis in 1:size(x,1)
                Evalπₓ[basis,cell,:] = temp[basis,cell]
            end
        end
        return Evalπₓ
    end

    # CDF method for scalar x-values
    function Πₓ(x::Real)
        out = zeros(n_phases(model))
        out[idx] = [zeros(1,sum(rates(model).>0)) pₓ] .+
        pₓ *
        [TDict["-+"]; TDict["0+"]] *
        (exp(K*x) - LinearAlgebra.I) / K *
        [LinearAlgebra.I(length(SDict["+"])) Ψ] *
        LinearAlgebra.diagm(1 ./ abs.(rates(model)[SDict["bullet"]])) *
        [LinearAlgebra.I(sum(rates(model) .!= 0)) [TDict["+0"];TDict["-0"]] * -T00inv]
        return out
    end
    # CDF method for arrays so that Πₓ returns an array with the same shape
    # as is output by Coeff2Dist
    function Πₓ(x::Array)
        temp = Πₓ.(x)
        Evalπₓ = zeros(Float64, size(x,1), size(x,2), n_phases(model))
        for cell in 1:size(x,2)
            for basis in 1:size(x,1)
                Evalπₓ[basis,cell,:] = temp[basis,cell]
            end
        end
        return Evalπₓ
    end

    return pₓ, πₓ, Πₓ, K
end

function stationary_distribution_x( model::Model)
    Ψ = psi_fun_x( model)
    ξ = xi_x( model, Ψ)
    pₓ, πₓ, Πₓ, K = stationary_distribution_x(model,Ψ,ξ)
    return pₓ, πₓ, Πₓ, K
end


include("../utils.jl")

@testset "Lorentz Manifold" begin
    M = Lorentz(3)
    @testset "Minkowski Metric" begin
        N = MetricManifold(Euclidean(3), MinkowskiMetric())
        @test N == M
        @test local_metric(M, zeros(3)) == Diagonal([1.0, 1.0, -1.0])
        # check minkowski metric is called
        p = zeros(3)
        X = [1.0, 2.0, 3.0]
        Y = [2.0, 3.0, 4.0]
        @test minkowski_metric(X, Y) == -4
        @test inner(M, p, X, Y) == minkowski_metric(X, Y)
    end
end


# PREAMBLE

# PKG_SETUP

# ## Setup

using ContactImplicitMPC# set environment variables to location of Gurobi software and license
ENV["GUROBI_HOME"] = "/gurobi912/linux64/"
ENV["GRB_LICENSE_FILE"] = "/gurobi912/gurobi.lic"

using Gurobi
using JuMP
using LinearAlgebra
using MeshCat
using Plots

include("/methods/structures.jl")
include("/methods/visuals.jl")

# ## Probem data
T = 40 # horizon (T+1 states, T controls)
H = 380 # simulation horizon
dt = 0.04 # step size
x0 = [0.00, 0.0] # initial state
Q = 1.0 # cost on state
Qf = 50.0 # cost on state
R = 1.0 # cost on control
β = 1e3 # extreme values in MIQP

# ## Dynamics model
model = WallPendulum(2, 1, 1.0, 1.0, 10.0, 1e4, 0.1) # m l g k d

A1, B1, c1 = dynamics_model(model, dt, mode = :none)
A2, B2, c2 = dynamics_model(model, dt, mode = :left)
A3, B3, c3 = dynamics_model(model, dt, mode = :right)

A = [A1, A2, A3]
B = [B1, B2, B3]
c = [c1, c2, c3]

# ## Domains
C1 = domain(model, mode = :none)
C2 = domain(model, mode = :left)
C3 = domain(model, mode = :right)
C = [C1, C2, C3]

# ## Problem
prob = WallProblem(model, T, x0, Q, Qf, R, β, A, B, c, C)

# ## Disturbances
i1 = 20
i2 = i1 + 100
i3 = i2 + 40
i4 = i3 + 100
i5 = i4 + 60
ind = [i1, i2, i3, i4, i5]
w = [[-15.5], [+15.5], [+15.5], [-11.5], [-10.5]]

# ## Simulation
x, u, s = simulate!(prob, x0, H, w = w, iw = ind)
sum(s) / (H * dt) # 4.52
maximum(s) / dt #17.6
mean(s) # 0.18
sqrt(mean((s .- mean(s)).^2)) #

# ##
DTu = range(0.0, step = dt, length = size(u)[1])
plot(DTu, hcat(u...)', label = "u")

# ##
DTx = range(0.0, step = dt, length = size(x)[1])
plot(DTx[1:end], hcat(x...)'[1:end,:], label = ["θ" "θd"])

# ## Visualizer
vis = Visualizer()
open(vis)

# ## Visualise
build_robot!(vis, model)
set_robot!(vis, model, [0.10, 0.1])
anim = visualize_robot!(vis, model, x, h = dt)
pθ_right = generate_pusher_traj(ind, w, x, side=:right)
pθ_left  = generate_pusher_traj(ind, w, x, side=:left)
visualize_disturbance!(vis, model, pθ_right, anim=anim, sample=1, offset=0.02, name=:PusherRight)
visualize_disturbance!(vis, model, pθ_left,  anim=anim, sample=1, offset=0.02, name=:PusherLeft)


using AtomicOpt
using LinearAlgebra
using Random
using Test

import AtomicOpt: LBFGSBQuad, Quadratic, boxquad, grad

"""Project into the box [bl, bu]."""
proj(x, bl, bu) = max.(min.(x, bu), bl)

@testset "Box-constrained LS via projected Newton" begin
    m, n = 10, 3
    A = randn(m, n); x0 = randn(n); b = A*x0
    bl = zeros(n); bu = ones(n)
    x = randn(n)

    obj(x) = 0.5*norm(A*x)^2 - dot(A'b,x)

    Q = Quadratic(A'*A, A'b)
    @test obj(x) ≈ Q(x)

end

@testset "Box-constrained LS via LBFGS-B" begin
    m, n = 100, 30
    A = randn(m, n); x0 = randn(n); b = A*x0
    bl = rand(n) .- 0.5; bu = rand(n) .+ 0.5

    Q = Quadratic(A'*A, A'b)
    lbfgsq = LBFGSBQuad(Q, bl, bu)
    x = boxquad(lbfgsq, factr=1.0)
    @test norm( x - proj(x-grad(Q,x), bl, bu)) ≤ 1e-6*norm(x)
end


export ExponentialLinearQuadratic

import Distributions: pdf, logpdf, ContinuousUnivariateDistribution

struct ExponentialLinearQuadratic{ A <: AbstractApproximationMethod, T <: Real } <: ContinuousUnivariateDistribution
    approximation :: A
    a :: T
    b :: T
    c :: T
    d :: T
end

ExponentialLinearQuadratic(approximation, a::Real, b::Real, c::Real, d::Real)             = ExponentialLinearQuadratic(approximation, promote(a, b, c, d)...)
ExponentialLinearQuadratic(approximation, a::Integer, b::Integer, c::Integer, d::Integer) = ExponentialLinearQuadratic(approximation, float(a), float(b), float(c), float(d))

Distributions.pdf(dist::ExponentialLinearQuadratic, x::Real)    = exp(logpdf(dist, x))
Distributions.logpdf(dist::ExponentialLinearQuadratic, x::Real) = -0.5 * (dist.a * x + dist.b * exp(dist.c * x + dist.d * x ^ 2 / 2.0))
Distributions.mean(dist::ExponentialLinearQuadratic)            = approximate_meancov(dist.approximation, (z) -> pdf(dist, z) * exp(0.5 * z^2), NormalMeanVariance(0.0,1.0))[1]
Distributions.var(dist::ExponentialLinearQuadratic)             = approximate_meancov(dist.approximation, (z) -> pdf(dist, z) * exp(0.5 * z^2), NormalMeanVariance(0.0,1.0))[2]
Distributions.cov(dist::ExponentialLinearQuadratic)             = var(dist)

precision(dist::ExponentialLinearQuadratic)    = inv(var(dist))
weightedmean(dist::ExponentialLinearQuadratic) = precision(dist)*mean(dist)

prod_analytical_rule(::Type{ <: UnivariateNormalDistributionsFamily }, ::Type{ <: ExponentialLinearQuadratic }) = ProdAnalyticalRuleAvailable()

function prod(::ProdAnalytical, left::UnivariateNormalDistributionsFamily, right::ExponentialLinearQuadratic)
    mean, variance = approximate_meancov(right.approximation, (z) -> pdf(right, z), left)
    return NormalMeanVariance(mean, variance)
end


using LinearAlgebra,DataFrames,Plots,StatsBase

function ManhattanPlot(inputFile,alpha = 0.1)
    @warn("This plot is using backend:",backend())

    maxPos_CHR = sort(by(inputFile, :chromosome, chr_len = :position => maximum),[:chromosome])
    maxPos_CHR[:tot] = cumsum(maxPos_CHR[:chr_len]) - maxPos_CHR[:chr_len]
    deletecols!(maxPos_CHR, :chr_len)
    dat = join(inputFile, maxPos_CHR, on = :chromosome, kind = :left) # join both data sets
    sort!(dat, [:chromosome]); # sort data frame by :chromosome column
    dat[:BPcum] = dat[:position] + dat[:tot] # position for markers at the xaxis
    plot_dat = dat[[:chromosome, :pvalue, :BPcum]]
    plot_dat[:pvalue] = -log10.(plot_dat[:pvalue])

    #convert factors of :chromosome to columns
    plot_data = unstack(plot_dat, :chromosome, :pvalue)
    ncolumn = ncol(plot_data)

    # replace all missing values by NaN
    for i = 2:ncolumn
        plot_data[i] = coalesce.(plot_data[i], NaN)

    end

    plot_data = convert(Matrix, plot_data)
    axisdf = by(dat, :chromosome, center = :BPcum => x -> (maximum(x) + minimum(x))/2) # position for x axis
    p1 = scatter(plot_data[:,1], plot_data[:, 2:end], legend = false, xticks = (axisdf[:center], axisdf[:chromosome]),
    xlabel = "chromosome", ylab = ("-log(p)"))

    #Plot the significant line
    plot!([-log10(alpha)],seriestype = :hline,linecolor = "black")

    #display(p1)
end

#Testing
nmarkers  = 500
chrlength = 1*100*1_000_000 #one morgan
input1 = DataFrame(markerID = fill("m1",nmarkers), chromosome = sample(1:10,nmarkers),position = sample(1:chrlength,nmarkers), pvalue = rand(nmarkers))

#@time ManhattanPlot(input1,0.05)
ManhattanPlot(input1,0.05)


using Oceananigans.Architectures
using Oceananigans.Architectures: device
import Oceananigans.Architectures: architecture
using CUDA, CUDA.CUSPARSE
using KernelAbstractions: @kernel, @index

using LinearAlgebra, SparseArrays, IncompleteLU
using SparseArrays: nnz

import LinearAlgebra.ldiv!

"""
`ILUFactorization` (`preconditioner_method = :ILUFactorization`)
    stores two sparse lower and upper trianguilar matrices `L` and `U` such that `LU ≈ A`
    is applied to `r` with `forward_substitution!(L, r)` followed by `backward_substitution!(U, r)`
    constructed with `ilu(A, τ = drop_tolerance)`

`SparseInversePreconditioner` (`preconditioner_method = :SparseInverse` and `:AsymptoticInverse`)
    stores a sparse matrix `M` such that `M ≈ A⁻¹`
    is applied to `r` with a matrix multiplication `M * r`
    constructed with
    `asymptotic_diagonal_inverse_preconditioner(A)`
        -> is an asymptotic expansion of the inverse of A assuming that A is diagonally dominant
        -> it is possible to choose order 0 (Jacobi), 1 or 2
    `sparse_approximate_preconditioner(A, ε = tolerance, nzrel = relative_maximum_number_of_elements)`
        -> same formulation as Grote M. J. & Huckle T, "Parallel Preconditioning with sparse approximate inverses"
        -> starts constructing the sparse inverse of A from identity matrix until, either a tolerance (ε) is met or nnz(M) = nzrel * nnz(A)

The suggested preconditioners are

on the `CPU`
`ilu()` (superior to everything always and in every situation!)

on the `GPU`
`aymptotic_diagonal_inverse_preconditioner()` (if `Δt` is variable or large problem_sizes)
`sparse_inverse_preconditioner()` (if `Δt` is constant and problem_size is not too large)

as a rule of thumb, for poisson solvers:
`sparse_inverse_preconditioner` is better performing than `asymptotic_diagonal_inverse_preconditioner` only if `nzrel >= 2.0`
As such, we urge to use `sparse_inverse_preconditioner` only when
- Δt is constant (we don't have to recalculate the preconditioner during the simulation)
- it is feasible to choose `nzrel = 2.0` (for not too large problem sizes)

Note that `asymptotic_diagonal_inverse_preconditioner` assumes the matrix to be diagonally dominant, for this reason it could
be detrimental when used on non-diagonally dominant system (cases where Δt is very large). In this case it is better
to use `sparse_inverse_preconditioner`

`ilu()` cannot be used on the GPU because preconditioning the solver with a direct LU (or Choleski) type
of preconditioner would require too much computation for the `ldiv!(P, r)` step completely hindering the performances
"""

validate_settings(T, arch, settings)                                  = settings
validate_settings(::Val{:Default}, arch, settings)                    = arch isa CPU ? (τ = 0.001, ) : (order = 1, )
validate_settings(::Val{:SparseInverse}, arch, settings::Nothing)     = (ε = 0.1, nzrel = 2.0)
validate_settings(::Val{:ILUFactorization}, arch, settings::Nothing)  = (τ = 0.001, )
validate_settings(::Val{:AsymptoticInverse}, arch, settings::Nothing) = (order = 1, )

validate_settings(::Val{:ILUFactorization}, arch, settings)  = haskey(settings, :τ) ?
                                                                      settings :
                                                                      throw(ArgumentError("τ has to be specified for ILUFactorization"))
validate_settings(::Val{:SparseInverse}, arch, settings)     = haskey(settings, :ε) && haskey(settings, :nzrel) ?
                                                                      settings :
                                                                      throw(ArgumentError("both ε and nzrel have to be specified for SparseInverse"))
validate_settings(::Val{:AsymptoticInverse}, arch, settings) = haskey(settings, :order) ?
                                                                      settings :
                                                                      throw(ArgumentError("and order ∈ [0, 1, 2] has to be specified for AsymptoticInverse"))

function build_preconditioner(::Val{:Default}, matrix, settings)
    default_method = architecture(matrix) isa CPU ? :ILUFactorization : :AsymptoticInverse
    return build_preconditioner(Val(default_method), matrix, settings)
end

build_preconditioner(::Val{nothing},            A, settings)  = Identity()
build_preconditioner(::Val{:SparseInverse},     A, settings)  = sparse_inverse_preconditioner(A, ε = settings.ε, nzrel = settings.nzrel)
build_preconditioner(::Val{:AsymptoticInverse}, A, settings)  = asymptotic_diagonal_inverse_preconditioner(A, asymptotic_order = settings.order)

function build_preconditioner(::Val{:ILUFactorization},  A, settings)
    if architecture(A) isa GPU
        throw(ArgumentError("the ILU factorization is not available on the GPU! choose another method"))
    else
        return ilu(A, τ = settings.τ)
    end
end

@inline architecture(::CuSparseMatrixCSC) = CUDAGPU()
@inline architecture(::SparseMatrixCSC)   = CPU()

abstract type AbstractInversePreconditioner{M} end

function  LinearAlgebra.ldiv!(u, precon::AbstractInversePreconditioner, v)
    mul!(u, matrix(precon), v)
end

function  LinearAlgebra.ldiv!(precon::AbstractInversePreconditioner, v)
    mul!(v, matrix(precon), v)
end

struct SparseInversePreconditioner{M} <: AbstractInversePreconditioner{M}
    Minv :: M
end

@inline matrix(p::SparseInversePreconditioner)  = p.Minv

"""
The diagonally dominant inverse preconditioner is constructed with an asymptotic expansion of `A⁻¹` up to the second order
If `I` is the Identity matrix and `D` is the matrix containing the diagonal of `A`, then

the 0th order expansion is the Jacobi preconditioner `M = D⁻¹ ≈ A⁻¹`

the 1st order expansion corresponds to `M = D⁻¹(I - (A - D)D⁻¹) ≈ A⁻¹`

the 2nd order expansion corresponds to `M = D⁻¹(I - (A - D)D⁻¹ + (A - D)D⁻¹(A - D)D⁻¹) ≈ A⁻¹`

all preconditioners are calculated on CPU and then moved to the GPU.
Additionally the first order expansion has a method to calculate the preconditioner directly on the GPU
`asymptotic_diagonal_inverse_preconditioner_first_order(A)` in case of variable time step where the preconditioner
has to be recalculated often
"""

function asymptotic_diagonal_inverse_preconditioner(A::AbstractMatrix; asymptotic_order)

    arch                  = architecture(A)
    constr                = deepcopy(constructors(arch, A))
    colptr, rowval, nzval = copy_unpack_constructors(arch, constr)
    dev                   = device(arch)

    M = size(A, 1)

    invdiag = arch_array(arch, zeros(eltype(nzval), M))

    loop! = _get_inv_diag!(dev, 256, M)
    event = loop!(invdiag, colptr, rowval, nzval; dependencies=Event(dev))
    wait(dev, event)

    if asymptotic_order == 0
        Minv_cpu = spdiagm(0=>arch_array(CPU(), invdiag))
        Minv     = arch_sparse_matrix(arch, Minv_cpu)
    elseif asymptotic_order == 1
        loop! = _initialize_asymptotic_diagonal_inverse_preconditioner_first_order!(dev, 256, M)
        event = loop!(nzval, colptr, rowval, invdiag; dependencies=Event(dev))
        wait(dev, event)

        constr_new = (colptr, rowval, nzval)
        Minv = arch_sparse_matrix(arch, constructors(arch, M, constr_new))
    else
        D   = spdiagm(0=>diag(arch_sparse_matrix(CPU(), A)))
        D⁻¹ = spdiagm(0=>arch_array(CPU(), invdiag))
        Minv_cpu = D⁻¹ * (I - (A - D) * D⁻¹ + (A - D) * D⁻¹ * (A - D) * D⁻¹)
        Minv = arch_sparse_matrix(arch, Minv_cpu)
    end

    return SparseInversePreconditioner(Minv)
end

@kernel function _initialize_asymptotic_diagonal_inverse_preconditioner_first_order!(nzval, colptr, rowval, invdiag)
    col = @index(Global, Linear)

    for idx = colptr[col] : colptr[col+1] - 1
        if rowval[idx] == col
            nzval[idx] = invdiag[col]
        else
            nzval[idx] = - nzval[idx] * invdiag[rowval[idx]] * invdiag[col]
        end
    end
end

function sparse_inverse_preconditioner(A::AbstractMatrix; ε, nzrel)

   # let's choose an initial sparsity => diagonal
   A_cpu    = arch_sparse_matrix(CPU(), A)
   Minv_cpu = sparse_approximate_inverse(A_cpu, ε = ε, nzrel = nzrel)

   Minv = arch_sparse_matrix(architecture(A), Minv_cpu)
   return SparseInversePreconditioner(Minv)
end


using ApproxInferenceBase
using Test
using Distributions
using Random

@testset "Factored" begin
    d=Factored(Uniform(0,1),Uniform(100,101))
    @test all((0,100) .<= rand(d) .<= (1,101))
    @test pdf(d,(0.0,0.0)) == 0.0
    @test pdf(d,(0.5,100.5)) == 1.0
    @test logpdf(d,(0.5,100.5)) == 0.0
    @test logpdf(d,(0.0,0.0)) == -Inf
    @test length(d) == 2
    m=Factored(Uniform(0.00,1.0),DiscreteUniform(1,2))
    sample=rand(m)
    @test 0<sample[1]<1
    @test sample[2] == 1 || sample[2] == 2
    @test pdf(m,sample) == 0.5
    @test logpdf(m,sample) ≈ log(0.5)
end


using FreudenthalTriangulation
using Test

@testset "FreudenthalTriangulation.jl" begin
    # test freudenthal_vertices
    n, m = rand(1:10, 2)
    V = freudenthal_vertices(n, m)
    @test size(V, 1) == binomial(n + m - 1, m)
    v = rand(V)
    @test size(v, 1) == n && v[1] == m
    @test issorted(v, rev=true)

    # test freudenthal_simplex
    n = rand(5:20)
    x = rand(Float64, n)
    V = freudenthal_simplex(x)
    @test size(V, 1) == n + 1
    x = rand(Int64, n)
    x = convert(Vector{Float64}, x)
    V = freudenthal_simplex(x)
    distances = [abs(sum(x - V[i])) for i in 1:n]
    @test issorted(distances)

    # test barycentric_coordinates
    n = rand(5:20)
    x = rand(Float64, n)
    V = freudenthal_simplex(x)
    λ = barycentric_coordinates(x, V)
    @test size(λ, 1) == n + 1
    @test abs(sum(λ) - 1.0) < 1e-8

    # test freudenthal_simplex_and_coords!
    n = rand(5:20)
    x = rand(Float64, n)
    V = freudenthal_simplex(x)
    λ = barycentric_coordinates(x, V)
    V2, λ2 = freudenthal_simplex_and_coords(x)
    i = rand(1:n)
    v = V[i]
    v2 = V2[i]
    @test abs(sum(v - v2)) < 1e-8
    @test abs(sum(λ - λ2)) < 1e-8

    # test freudenthal_simplex_and_coords!
    n = rand(5:20)
    x = rand(Float64, n)
    V = freudenthal_simplex(x)
    λ = barycentric_coordinates(x, V)
    V2 = [rand(Int64, n) for i=1:n+1]
    λ2 = rand(Float64, n+1)
    V2, λ2 = freudenthal_simplex_and_coords!(x, V2, λ2)
    i = rand(1:n)
    v = V[i]
    v2 = V2[i]
    @test abs(sum(v - v2)) < 1e-8
    @test abs(sum(λ - λ2)) < 1e-8
end


include("activation_functions.jl")

"""
Make a linear forward calc.
"""
function linear_forward(A, W, b)
    # Make a linear forward and return the inputs as cache
    Z = BigFloat.((W * A) .+ b)
    cache = (A, W, b)

    @assert size(Z) == (size(W, 1), size(A, 2))

    return (Z = Z, cache = cache)
end # function

"""
Make a forward activation
"""
function linear_forward_activate(A_prev, W, b; activation_function = "tanh")
    @assert activation_function ∈ ("sigmoid", "relu", "softmax", "tanh", "swish")
    Z, linear_cache = linear_forward(A_prev, W, b)

    if activation_function == "sigmoid"
        A, activation_cache = sigmoid(Z)
    elseif activation_function == "relu"
        A, activation_cache = relu(Z)
    elseif activation_function == "softmax"
        A, activation_cache = softmax(Z)
    elseif activation_function == "tanh"
        A, activation_cache = tanhact(Z)
    elseif activation_function == "swish"
        A, activation_cache = swish(Z)
    end # if

    cache = (linear_step_cache = linear_cache, activation_step_cache = activation_cache)

    @assert size(A)[1] == size(W, 1)[1]

    return A, cache
end # function


function calculate_u1u2_spectrum!(Ef,u,cplane)
    ux = u.c.x
    uy = u.c.y
    # Initialize the shells to zeros
    fill!(Ef,0)
    maxdk2d = max(KX[2],KY[2])
    nshells = min(NX,NY÷2)
    @inbounds for j in YRANGE
        conjFactX=1.0
        for i in XRANGE
            k = fsqrt(KX[i]^2 + KY[j]^2)
            n = round(Int, k/maxdk2d) + 1
            if n <= nshells
                magsq = abs2(ux[i,j,cplane]) + abs2(uy[i,j,cplane])
                ee = 0.5*conjFactX * magsq;
                Ef[n]+=ee;
            end
            conjFactX=2.0;
        end
    end
end

function calculate_u1u2_spectrum(u,cplane,s::AbstractSimulation)
    Ef = zeros(min(NX, NY÷2))
    calculate_u1u2_spectrum!(Ef,u,1)
    return Ef
end

function calculate_abs2!(Ef,u,cplane)
    ux = u.c.x
    uy = u.c.y
    # Initialize the shells to zeros
    fill!(Ef,0)
    @inbounds for j in YRANGE
        conjFactX=1.0
        for i in XRANGE
            magsq = abs2(ux[i,j,cplane]) + abs2(uy[i,j,cplane])
            ee = 0.5*conjFactX * magsq;
            Ef[i,j]=ee;
            conjFactX=2.0;
        end
    end
end

function calculate_u1u2_mag_spectrum!(Ef,u,cplane)
    ux = u.c.x
    uy = u.c.y
    # Initialize the shells to zeros
    fill!(Ef,0)
    maxdk2d = max(KX[2],KY[2])
    nshells = min(NX,NY÷2)
    @inbounds for j in YRANGE
        conjFactX=1.0
        for i in XRANGE
            k = fsqrt(KX[i]^2 + KY[j]^2)
            n = round(Int, k/maxdk2d) + 1
            if n <= nshells
                mag = fsqrt(abs2(ux[i,j,cplane]) + abs2(uy[i,j,cplane]))
                ee = conjFactX * mag;
                Ef[n]+=ee;
            end
            conjFactX=2.0;
        end
    end
end


module LogisticRegression1

using Plots
using Distributions
using LinearAlgebra

# Generate dummy data
n = 100;
x1 = transpose(rand(MvNormal([1., 2.], I), n));
x2 = transpose(rand(MvNormal([-2., -1.5], I), n));
y1 = repeat([1], n);
y2 = repeat([-1], n);
p = scatter(x1[:, 1], x1[:, 2], color = :blue)
scatter!(p, x2[:, 1], x2[:, 2], color = :red)

using Flux
using Flux.Tracker: gradient, update!, data
using Random
using Formatting: printfmt

# Logistic regression
Random.seed!(666);
α = param(rand());   # Bias
β = param(rand(2, 1));  # Weight
predict(X) = σ.(α .+ X*β)
X = [x1; x2];
y = [y1; y2];

function plot_contour(xlim, ylim, f)
    x = xlim[1]:0.5:xlim[2]
    y = ylim[1]:0.5:ylim[2]
    X = repeat(reshape(x, 1, :), length(y), 1)
    Y = repeat(y, 1, length(x))
    contour(x, y, f, fill=true, color=:bluesreds)
    scatter!(x1[:, 1], x1[:, 2],
            color = :red, label="Class 1")
    scatter!(x2[:, 1], x2[:, 2],
            color = :blue, label="Class 2")
end

xlim = ylim = (-5, 5)``
plot_contour(xlim, ylim,
                 (x, y) -> (first ∘ data ∘ predict)([x y]))

function loss(x, y)
    ŷ = predict(x)
    mean(Flux.binarycrossentropy.(ŷ, y))
end

printfmt("Initial loss: {:.4f}", (data ∘ loss)(X, y));

η = 0.01 # learning rate
epochs = 200
plot_interval = 10
anim = @animate for i=1:epochs
    for i in 1:plot_interval
        gs = gradient(() -> loss(X, y), Flux.params(α, β))
        update!(β, -η*gs[β])
        #update!(α, -η*gs[α])
    end
    display(plot_contour(xlim, ylim, (x, y) ->
                 (first ∘ data ∘ predict)([x y])))
end
gif(anim, "bluesreds.gif", fps = 60)

printfmt("Loss after training: {:.4f}", (data ∘ loss)(X, y));

end


#
# Vadd
#
# Element wise addition of two vectors at a time in a chain (C=A+B; D=C+E; F=D+G)
# Asks the user to select a device at runtime
#
# History: Initial version based on vadd.c, written by Tim Mattson, June 2011
# Ported to C++ Wrapper API by Benedict Gaster, September 2011
# Updated to C++ Wrapper API v1.2 by Tom Deakin and Simon McIntosh-Smith, October 2012
# Ported to Python by Tom Deakin, July 2013
# Ported to Julia  by Jake Bolewski, Nov 2013

using OpenCL

# tolerance used in floating point comparisons
TOL = 1e-3

# length of vectors a, b, c
LENGTH = 1024

# Kernel: vadd
#
# To compute the elementwise sum c = a + b
#
# Input: a and b float vectors of length count
# Output c float vector of length count holding the sum a + b

kernelsource = "
__kernel void vadd(
    __global float* a,
    __global float* b,
    __global float* c,
    const unsigned int count)
{
    unsigned int i = get_global_id(0);
    if (i < count)
        c[i] = a[i] + b[i];
}
"

# create a compute context

# this selects the fastest opencl device available
# and creates a context and queue for using the
# the selected device
device, ctx, queue = cl.create_compute_context()

# create the compute program and build it
program = cl.Program(ctx, source=kernelsource) |> cl.build!

#create a, b, e, and g vectors and fill with random float values
#create empty vectors for c, d, and f
h_a = rand(Float32, LENGTH)
h_b = rand(Float32, LENGTH)
h_c = Vector{Float32}(undef, LENGTH)
h_d = Vector{Float32}(undef, LENGTH)
h_e = rand(Float32, LENGTH)
h_f = Vector{Float32}(undef, LENGTH)
h_g = rand(Float32, LENGTH)

# create the input (a,b,e,g) arrays in device memory and copy data from the host

# buffers can be passed memory flags:
# {:r = readonly, :w = writeonly, :rw = read_write (default)}

# buffers can also be passed flags for allocation:
# {:use (use host buffer), :alloc (alloc pinned memory), :copy (default)}

# Create the input (a, b, e, g) arrays in device memory and copy data from host
d_a = cl.Buffer(Float32, ctx, (:r, :copy), hostbuf=h_a)
d_b = cl.Buffer(Float32, ctx, (:r, :copy), hostbuf=h_b)
d_e = cl.Buffer(Float32, ctx, (:r, :copy), hostbuf=h_e)
d_g = cl.Buffer(Float32, ctx, (:r, :copy), hostbuf=h_g)
# Create the output (c, d, f) array in device memory
d_c = cl.Buffer(Float32, ctx, :w, LENGTH)
d_d = cl.Buffer(Float32, ctx, :w, LENGTH)
d_f = cl.Buffer(Float32, ctx, :w, LENGTH)

# create the kernel
vadd = cl.Kernel(program, "vadd")

# execute the kernel over the entire range of 1d input
# calling `queue` is blocking, it accepts the kernel, global / local work sizes,
# the the kernel's arguments.

# here we call the kernel with work size set to the number of elements and a local
# work size of nothing. This enables the opencl runtime to optimize the local size
# for simple kernels
queue(vadd, size(h_a), nothing, d_a, d_b, d_c, UInt32(LENGTH))

# an alternative syntax is to create an partial function to call
# by julia's getindex syntax for Kernel types.
# here the queue, global_size, and (optional) local_size are passed in which
# returns a partial function with these parameters set.
vadd[queue, size(h_e)](d_e, d_c, d_d, UInt32(LENGTH))
vadd[queue, size(h_g)](d_g, d_d, d_f, UInt32(LENGTH))

# copy back the results from the compute device
# copy!(queue, dst, src) follows same interface as julia's built in copy!
cl.copy!(queue, h_f, d_f)

# test the results
correct = 0
for i in 1:LENGTH
    tmp = h_a[i] + h_b[i] + h_e[i] + h_g[i]
    tmp -= h_f[i]
    if tmp^2 < TOL^2
        global correct += 1
    else
        println("tmp $tmp h_a $(h_a[i]) h_b $(h_b[i]) ",
                "h_e $(h_e[i]) h_g $(h_g[i]) h_f $(h_f[i])")
    end
end

# summarize results
println("3 vector adds to find F=A+B+E+G: $correct out of $LENGTH results were correct")


# # Stability

# In this example, we will demonstrate the stability of persistent homology. The stability
# theorem roughly states that a small change in the input data will result in a small
# change in the resulting persistence diagram. In other words, persistent homology is very
# tolerant of noisy data. Also, see the
# [Distances](https://mtsch.github.io/PersistenceDiagrams.jl/dev/generated/distances/)
# example in [PersistenceDiagrams.jl](https://github.com/mtsch/PersistenceDiagrams.jl).

# Again, start with loading some packages.

using Ripserer
using Plots
using Random; # hide
Random.seed!(1337); # hide
gr(); # hide
nothing; # hide

# As in the Basics example, we will look at the persistent homology of a noisy circle.

function noisy_circle(n; r1=1, r2=1, noise=0.1)
    points = NTuple{2,Float64}[]
    for _ in 1:n
        θ = 2π * rand()
        point = (
            r1 * sin(θ) + noise * rand() - noise / 2,
            r2 * cos(θ) + noise * rand() - noise / 2,
        )
        push!(points, point)
    end
    return points
end

# We will look at the first persistent homology group of this space.

# First, let's see what happens if we repeatedly sample 100 random points from a circle.

anim = @animate for _ in 1:200
    points = noisy_circle(100; noise=0)
    result = ripserer(points)

    plt_pts = scatter(
        points;
        legend=false,
        aspect_ratio=1,
        xlim=(-2.2, 2.2),
        ylim=(-2.2, 2.2),
        title="Data",
    )
    plt_diag = plot(result; infinity=3)

    plot(plt_pts, plt_diag; size=(800, 400))
end
gif(anim, "stability_anim_1.gif") # hide

# We notice that an interval in ``H_1`` always stands out and that its death remains
# constant. The only thing that changes is the birth time. The birth time is equal to the
# largest distance between adjacent points in the circle. At the birth time, the circle is
# connected.

# Now, let's add some noise!

anim = @animate for _ in 1:200
    points = noisy_circle(100; noise=0.2)
    result = ripserer(points)

    plt_pts = scatter(
        points;
        legend=false,
        aspect_ratio=1,
        xlim=(-2.2, 2.2),
        ylim=(-2.2, 2.2),
        title="Data",
    )
    plt_diag = plot(result; infinity=3)

    plot(plt_pts, plt_diag; size=(800, 400))
end
gif(anim, "stability_anim_2.gif") # hide

# The interval is jumping around a lot more now, but it hovers around the same general area.
# It's still clearly the most persistent feature of our space.

# Next, let's look at how adding more and more noise affects the diagram.

anim = @animate for noise in vcat(0:0.01:1, 1:-0.01:0)
    points = noisy_circle(100; noise=noise)
    result = ripserer(points)

    plt_pts = scatter(
        points;
        legend=false,
        aspect_ratio=1,
        xlim=(-2.2, 2.2),
        ylim=(-2.2, 2.2),
        title="Data",
    )
    plt_diag = plot(result; infinity=3)

    plot(plt_pts, plt_diag; size=(800, 400))
end
gif(anim, "stability_anim_3.gif") # hide

# We see we have to add quite a bit of noise to destroy the diagram. Notice how the death
# time of the interval decreases as we add noise. This is the result of the diameter of the
# hole in our circle shrinking.

# Finally, let's stretch our circle.

anim = @animate for r in vcat(0.0:0.02:2, 2:-0.02:0.0)
    points = noisy_circle(100; noise=0.1, r1=r)
    result = ripserer(points)

    plt_pts = scatter(
        points;
        legend=false,
        aspect_ratio=1,
        xlim=(-2.2, 2.2),
        ylim=(-2.2, 2.2),
        title="Data",
    )
    plt_diag = plot(result; infinity=3)

    plot(plt_pts, plt_diag; size=(800, 400))
end
gif(anim, "stability_anim_4.gif") # hide

# Again, we see the persistent homology stays stable, as long as the data at least somewhat
# resembles a circle.


# %% show AC, MI for choosing optimal τ
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot, Random

ττ = 0:90
ac = autocor(w, ττ)
smi = selfmutualinfo(w, ττ)
smi ./= maximum(smi)

fig = figure(figsize = (figx/2, figy))
ax = gca()
ax.plot(ττ, ac; label = "AC")
ax.plot(ττ, smi, label = "SMI", color = "C3")
τ = estimate_delay(w, "mi_min")
ax.scatter(τ, smi[τ]; color = "C2", s = 100, zorder = 99)
ax.set_ylim(0,1.05)
ax.set_xticks(0:30:90)
ax.set_xlabel("\$\\tau\$", labelpad = -20)
ax.legend()
fig.tight_layout(;pad = 0.25)
wsave(plotsdir("6", "mutualinfo_ac_tau"), fig)


# Tests for Multinomial

using Distributions
using Base.Test

p = [0.2, 0.5, 0.3]
nt = 10
d = Multinomial(nt, p)

# Basics

@test length(d) == 3
@test d.n == nt
@test mean(d) ≈ [2., 5., 3.]
@test var(d)  ≈ [1.6, 2.5, 2.1]
@test cov(d)  ≈ [1.6 -1.0 -0.6; -1.0 2.5 -1.5; -0.6 -1.5 2.1]

@test insupport(d, [1, 6, 3])
@test !insupport(d, [2, 6, 3])
@test partype(d) == Float64
@test partype(Multinomial(nt, Vector{Float32}(p))) == Float32

# Conversion
@test typeof(d) == Multinomial{Float64}
@test typeof(Multinomial(nt, Vector{Float32}(p))) == Multinomial{Float32}
@test typeof(convert(Multinomial{Float32}, d)) == Multinomial{Float32}
@test typeof(convert(Multinomial{Float32}, params(d)...)) == Multinomial{Float32}

# random sampling

x = rand(d)
@test isa(x, Vector{Int})
@test sum(x) == nt
@test insupport(d, x)
@test size(x) == size(d)
@test length(x) == length(d)
@test d == typeof(d)(params(d)...)

x = rand(d, 100)
@test isa(x, Matrix{Int})
@test all(sum(x, 1) .== nt)
@test all(insupport(d, x))

x = rand(sampler(d))
@test isa(x, Vector{Int})
@test sum(x) == nt
@test insupport(d, x)

# logpdf

x1 = [1, 6, 3]

@test isapprox(pdf(d, x1), 0.070875, atol=1.0e-8)
@test logpdf(d, x1) ≈ log(pdf(d, x1))

x = rand(d, 100)
pv = pdf(d, x)
lp = logpdf(d, x)
for i in 1 : size(x, 2)
	@test pv[i] ≈ pdf(d, x[:,i])
	@test lp[i] ≈ logpdf(d, x[:,i])
end

# test type stability of logpdf
@test typeof(logpdf(convert(Multinomial{Float32}, d), x1)) == Float32

# suffstats

d0 = d
n0 = 100
x = rand(d0, n0)
w = rand(n0)

ss = suffstats(Multinomial, x)
@test isa(ss, Distributions.MultinomialStats)
@test ss.n == nt
@test ss.scnts == vec(sum(Float64[x[i,j] for i = 1:size(x,1), j = 1:size(x,2)], 2))
@test ss.tw == n0

ss = suffstats(Multinomial, x, w)
@test isa(ss, Distributions.MultinomialStats)
@test ss.n == nt
@test ss.scnts ≈ Float64[x[i,j] for i = 1:size(x,1), j = 1:size(x,2)] * w
@test ss.tw    ≈ sum(w)

# fit

x = rand(d0, 10^5)
@test size(x) == (length(d0), 10^5)
@test all(sum(x, 1) .== nt)

r = fit(Multinomial, x)
@test r.n == nt
@test length(r) == length(p)
@test isapprox(probs(r), p, atol=0.02)

r = fit_mle(Multinomial, x, fill(2.0, size(x,2)))
@test r.n == nt
@test length(r) == length(p)
@test isapprox(probs(r), p, atol=0.02)

# behavior for n = 0
d0 = Multinomial(0, p)
@test rand(d0) == [0, 0, 0]
@test pdf(d0, [0, 0, 0]) == 1
@test pdf(d0, [0, 1, 0]) == 0
@test mean(d0) == [0, 0, 0]
@test var(d0) == [0, 0, 0]
@test cov(d0) == zeros(3, 3)
@test entropy(d0) == 0
@test insupport(d0, [0, 0, 0]) == true
@test insupport(d0, [0, 0, 4]) == false
@test length(d0) == 3
@test size(d0) == (3,)


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

# %%
using LinearAlgebra
dot2(v) = dot(v, v)
using Plots

n = 9
z = exp(2π*im/n)
t = range(0, 2π; length=101)
v0 = normalize([z^k for k in 0:n-1])
#v0 = normalize!(rand(ComplexF64, n))
v(t) = exp(im*t)*v0
f(t) = real(dot2(real(v(t))))
g(t) = real(dot2(imag(v(t))))
plot(t, f.(t); label="f(t)", ylim=(0, 1))
plot!(t, g.(t); label="g(t)", ls=:auto)

# %%
using LinearAlgebra
dot2(v) = dot(v, v)
using Plots

n = 9
z = exp(2π*im/n)
t = range(0, 2π; length=101)
#v0 = normalize([z^k for k in 0:n-1])
v0 = normalize!(rand(ComplexF64, n))
v(t) = exp(im*t)*v0
f(t) = real(dot2(real(v(t))))
g(t) = real(dot2(imag(v(t))))
plot(t, f.(t); ylim=(0, 1))
plot!(t, g.(t); ls=:auto)

# %%
using LinearAlgebra
using Plots

n = 10
v0 = normalize(randn(ComplexF64, n))
v(t) = exp(im*t)*v0
f(t) = real(dot2(real(v(t))))
g(t) = real(dot2(imag(v(t))))

s0 = sum(z^2 for z in v0)
n0 = norm(v0)^2
F(t) = (n0 + real(exp(im*2t) * s0))/2
G(t) = (n0 - real(exp(im*2t) * s0))/2

t = range(0, 2π; length=201)
P = plot(t, f.(t); label="f(t)", lw=2)
plot!(t, F.(t); label="F(t)", ls=:auto, lw=2, ylim=(0, 1))
plot!(t, g.(t); label="g(t)", lw=2)
plot!(t, G.(t); label="G(t)", ls=:auto, lw=2, ylim=(0, 1))
plot!(; legend=:outertopright)

# %%
using LinearAlgebra
using Plots

n = 10
v0 = normalize(randn(ComplexF64, n))
v(t) = exp(im*t)*v0
f(t) = real(dot2(real(v(t))))
g(t) = real(dot2(imag(v(t))))

s0 = sum(z^2 for z in v0)
n0 = norm(v0)^2
F(t) = (n0 + real(exp(im*2t) * s0))/2
G(t) = (n0 - real(exp(im*2t) * s0))/2

t = range(0, 2π; length=201)
P = plot(t, f.(t); label="f(t)", lw=2)
plot!(t, F.(t); label="F(t)", ls=:auto, lw=2, ylim=(0, 1))
plot!(t, g.(t); label="g(t)", lw=2)
plot!(t, G.(t); label="G(t)", ls=:auto, lw=2, ylim=(0, 1))
plot!(; legend=:outertopright)

# %%
using LinearAlgebra
using Plots

n = 10
v0 = normalize(randn(ComplexF64, n))
v(t) = exp(im*t)*v0
f(t) = real(dot2(real(v(t))))
g(t) = real(dot2(imag(v(t))))

s0 = sum(z^2 for z in v0)
n0 = norm(v0)^2
F(t) = (n0 + real(exp(im*2t) * s0))/2
G(t) = (n0 - real(exp(im*2t) * s0))/2

t = range(0, 2π; length=201)
P = plot(t, f.(t); label="f(t)", lw=2)
plot!(t, F.(t); label="F(t)", ls=:auto, lw=2, ylim=(0, 1))
plot!(t, g.(t); label="g(t)", lw=2)
plot!(t, G.(t); label="G(t)", ls=:auto, lw=2, ylim=(0, 1))
plot!(; legend=:outertopright)

# %%
using LinearAlgebra
using Plots

n = 10
v0 = normalize(randn(ComplexF64, n))
v(t) = exp(im*t)*v0
f(t) = real(dot2(real(v(t))))
g(t) = real(dot2(imag(v(t))))

s0 = sum(z^2 for z in v0)
n0 = norm(v0)^2
F(t) = (n0 + real(exp(im*2t) * s0))/2
G(t) = (n0 - real(exp(im*2t) * s0))/2

t = range(0, 2π; length=201)
P = plot(t, f.(t); label="f(t)", lw=2)
plot!(t, F.(t); label="F(t)", ls=:auto, lw=2, ylim=(0, 1))
plot!(t, g.(t); label="g(t)", lw=2)
plot!(t, G.(t); label="G(t)", ls=:auto, lw=2, ylim=(0, 1))
plot!(; legend=:outertopright)

# %%
using LinearAlgebra
using Plots

n = 10
v0 = normalize(randn(ComplexF64, n))
v(t) = exp(im*t)*v0
f(t) = real(dot2(real(v(t))))
g(t) = real(dot2(imag(v(t))))

s0 = sum(z^2 for z in v0)
n0 = norm(v0)^2
F(t) = (n0 + real(exp(im*2t) * s0))/2
G(t) = (n0 - real(exp(im*2t) * s0))/2

t = range(0, 2π; length=201)
P = plot(t, f.(t); label="f(t)", lw=2)
plot!(t, F.(t); label="F(t)", ls=:auto, lw=2, ylim=(0, 1))
plot!(t, g.(t); label="g(t)", lw=2)
plot!(t, G.(t); label="G(t)", ls=:auto, lw=2, ylim=(0, 1))
plot!(; legend=:outertopright)

# %%

0009/squared norm of the real part of a complex vector.jl


using Test
using CartesianMesh
#using Revise
using CutCell
include("useful_routines.jl")

cm = CutCell.CellMap([0.,2.],[3.,3.])
@test CutCell.dimension(cm) == 2
@test CutCell.jacobian(cm) == [1.5,0.5]
@test allapprox(cm([0.,0.]),[1.5,2.5])
@test allapprox(cm([1.,0.]),[3.,2.5])

mesh = UniformMesh([0.,0.],[1.,1.],[2,2])
@test allequal(CutCell.elements_per_mesh_side(mesh),[2,2])
cellmaps = CutCell.cell_maps(mesh)
cm1 = CutCell.CellMap([0.,0.],[0.5,0.5])
cm2 = CutCell.CellMap([0.,0.5],[0.5,1.])
cm3 = CutCell.CellMap([0.5,0.],[1.,0.5])
cm4 = CutCell.CellMap([0.5,0.5],[1.,1.])
@test allequal(cellmaps,[cm1,cm2,cm3,cm4])

cm = CutCell.CellMap([0.,0.],[1.,1.])
@test CutCell.determinant_jacobian(cm) ≈ 0.25

cm = CutCell.CellMap([1.,1.],[5.,8.])
@test CutCell.determinant_jacobian(cm) ≈ 7.0

cm = CutCell.CellMap([2.,1.],[4.,2.])
@test allapprox(CutCell.inverse(cm,[3,1.5]),[0.,0.])
@test allapprox(CutCell.inverse(cm,[4,2.]),[1.,1.])
@test allapprox(CutCell.inverse(cm,[2,1]),[-1.,-1.])


Random.seed!(0)

# facts("qraddcol") do

    # context("β = 0") do
        m = 100
        A = randn(m,0)
        R = Array{Float64}(undef, 0,0)
        r = Vector{Float64}(undef, m)

        for i in 1:m
            a = randn(m)
            Rout = Array{Float64}(undef, i, i)
            qraddcol!(A, R, a, Rout, r)
            global R = qraddcol(A, R, a)
            global A = [A a]
            @assert  norm( R'*R - A'*A ) < 1e-5
            @assert  norm( Rout'*Rout - A'*A ) < 1e-5 "$(norm( Rout'*Rout - A'*A )) is not less than 1e-5"
        end

        b = sum(A, dims=2)[:]
        (x, rtest) = csne(R, A, b)
        @assert norm(rtest) < 1e-5
        @assert norm(x .- 1) < 1e-5

        x = csne!(R, A, b, r)
        @assert norm(r) < 1e-5
        @assert norm(x .- 1) < 1e-5

    # end

    # context("β > 0") do
        m = 100
        A = randn(m,0)
        R = Array{Float64}(undef, 0,0)
        β = 0.1
        for i in 1:m
            a = randn(m)
            B = [A a]
            global R = qraddcol(A, R, a, β)
            @assert norm( R'*R - B'*B - β^2*I ) < 1e-5
        end
    # end

# end


struct PosteriorCredibleInterval{TP<:Prior,TD<:AbstractDesign,TI<:Integer,TR<:Real} <: IntervalEstimator
    prior::TP
    design::TD
    XX::Array{TI,2}
    bounds::Array{TR,2}
    alpha::TR
end

string(ci::PosteriorCredibleInterval) = @sprintf "PosteriorCredibleInterval<%s>" string(ci.prior)

function PosteriorCredibleInterval(prior::Prior, design::AbstractDesign, α::Real; epsilon = 1e-4)

    ϵ = epsilon
    XX = sample_space(design)
    nn = size(XX, 1)

    bounds = zeros(Float64, nn, 2)
    for i in 1:nn
        x1, x2    = XX[i,:]
        posterior = update(prior, x1 + x2, n(design, x1))
        phat      = (x1 + x2) / n(design, x1)
        quantile  = prob -> Roots.find_zero(p -> cdf(p, posterior) - prob, (ϵ, 1 - ϵ))
        l = try
                quantile(α)
            catch e
                isa(e, Roots.ConvergenceFailed) ? 0.0 : e
        end
        u = try
                quantile(1 - α)
            catch e
                isa(e, Roots.ConvergenceFailed) ? 1.0 : e
        end
        u <= l ? error("credible interval collapsed") : nothing
        bounds[i,:] .= [l, u]
    end
    return PosteriorCredibleInterval{typeof(prior),typeof(design),eltype(XX),Float64}(prior, design, XX, bounds, α)
end


"""
    function ct_kenya(t,minpoint_y;schoolseffect = 0.,
                    measures_x = 24.,minpoint_x = 55.,returnpoint_x = 254.,schoolsreopen = 319. )

Defines piece-wise linear contact rates.
"""
function ct_kenya(t,minpoint_y;schoolseffect = 0.,
                    measures_x = 24.,minpoint_x = 55.,returnpoint_x = 254.,schoolsreopen = 319. )
    ct = 1.0*(t<measures_x)
    ct += (1.0 - schoolseffect  + ((t-measures_x)*(minpoint_y  - 1.0)/(minpoint_x - measures_x)))*(t ≥ measures_x && t < minpoint_x)
    ct += (minpoint_y - schoolseffect + ((t-minpoint_x)*(1.0   - minpoint_y)/(returnpoint_x-minpoint_x) ))*(t ≥minpoint_x && t < returnpoint_x)
    ct += (1.0 - schoolseffect )*(t ≥ returnpoint_x && t<schoolsreopen)
    ct += 1.0*(t ≥ schoolsreopen)
end

"""
    function twoSEStransmission(du,u,p,t)

In-place creation of vector field for the transmission model
"""
function twoSEStransmission(du,u,p,t)
    #Get state
	S₁,E₁,I₁,R₁,W₁,C₁,S₂,E₂,I₂,R₂,W₂,C₂,C_sero₁,C_sero₂  = u
    #Get parameters
	R₀,α,γ,ϵ,σ,N₁,N₂,ω,schooleffect,ct_min1,ct_min2 = p
    #Adjust for schools being open or shut
    Reff = R₀*(t < 24 || t ≥ 319) + R₀*schooleffect*(t≥ 24 && t< 319)
    #Rate of visiting POIs for lower and higher SES groups
    ct1 = ct_kenya(t,ct_min1)
    ct2 = ct_kenya(t,ct_min2)
    #Force of infection on lower and higher SES individuals visiting POIs
    λ₁ = γ*Reff*(ct1*ϵ*I₁ + ct2*(1-ϵ)*I₂)
    λ₂ = γ*Reff*(ct1*(1-ϵ)*I₁ + ct2*ϵ*I₂)
    #Dynamics
	du[1] = -ct1*(S₁/N₁)*λ₁
	du[2] = ct1*((S₁+σ*W₁)/N₁)*λ₁ - α*E₁
	du[3] = α*E₁ - γ*I₁
	du[4] = γ*I₁ - ω*R₁
	du[5] = ω*R₁ - σ*ct1*(W₁/N₁)*λ₁
	du[6] = ct1*((S₁+σ*W₁)/N₁)*λ₁
	du[7] = -ct2*(S₂/N₂)*λ₂
	du[8] = ct2*((S₂+σ*W₂)/N₂)*λ₂ - α*E₂
	du[9] = α*E₂ - γ*I₂
	du[10] = γ*I₂ - ω*R₂
	du[11] = ω*R₂ - σ*ct2*(W₂/N₂)*λ₂
	du[12] = ct2*((S₂+σ*W₂)/N₂)*λ₂
    du[13] = ct1*(S₁/N₁)*λ₁
    du[14] = ct2*(S₂/N₂)*λ₂
end

"""
    function twoSEStransmission_with_lower_transmissibility(du,u,p,t)

In-place creation of vector field for the transmission model. Model where previously infected are less transmissible.
"""
function twoSEStransmission_with_lower_transmissibility(du,u,p,t)
    #Get state
	S₁,E₁,I₁,R₁,W₁,C₁,S₂,E₂,I₂,R₂,W₂,C₂,C_sero₁,C_sero₂,EW₁,IW₁,EW₂,IW₂ = u
    #Get parameters
	R₀,α,γ,ϵ,σ,N₁,N₂,ω,schooleffect,ct_min1,ct_min2,ι = p
    #Adjust for schools being open or shut
    Reff = R₀*(t < 24 || t ≥ 319) + R₀*schooleffect*(t≥ 24 && t< 319)
    #Rate of visiting POIs for lower and higher SES groups
    ct1 = ct_kenya(t,ct_min1)
    ct2 = ct_kenya(t,ct_min2)
    #Force of infection on lower and higher SES individuals visiting POIs
    λ₁ = γ*Reff*(ct1*ϵ*(I₁ + ι*IW₁) + ct2*(1-ϵ)*(I₂ + ι*IW₂))
    λ₂ = γ*Reff*(ct1*(1-ϵ)*(I₁ + ι*IW₁) + ct2*ϵ*(I₂ + ι*IW₂))
    #Dynamics
	du[1] = -ct1*(S₁/N₁)*λ₁
	du[2] = ct1*(S₁/N₁)*λ₁ - α*E₁
	du[3] = α*E₁ - γ*I₁
	du[4] = γ*I₁ + γ*IW₁ - ω*R₁
	du[5] = ω*R₁ - σ*ct1*(W₁/N₁)*λ₁
	du[6] = ct1*((S₁+σ*W₁)/N₁)*λ₁
	du[7] = -ct2*(S₂/N₂)*λ₂
	du[8] = ct2*(S₂/N₂)*λ₂ - α*E₂
	du[9] = α*E₂ - γ*I₂
	du[10] = γ*I₂ + γ*IW₂ - ω*R₂
	du[11] = ω*R₂ - σ*ct2*(W₂/N₂)*λ₂
	du[12] = ct2*((S₂+σ*W₂)/N₂)*λ₂
    du[13] = ct1*(S₁/N₁)*λ₁
    du[14] = ct2*(S₂/N₂)*λ₂
    du[15] = ct1*σ*(W₁/N₁)*λ₁ - α*EW₁
    du[16] = α*EW₁ - γ*IW₁
    du[17] = ct2*σ*(W₂/N₂)*λ₂ - α*EW₂
    du[18] = α*EW₂ - γ*IW₂
end

function build_prob()
    u0 = [4.3e6*0.7 - 100.,100.,0.,0.,0.,0.,
            4.3e6*0.3 - 100.,100.,0.,0.,0.,0.,
            0.0,0.0]
    f = ODEFunction(twoSEStransmission;syms = [:S₁,:E₁,:I₁,:R₁,:W₁,:C₁,:S₂,:E₂,:I₂,:R₂,:W₂,:C₂,:C_sero₁,:C_sero₂])
    return ODEProblem(f,u0,(0,365))
end

function build_prob(::Val{:lower_trans})
    u0 = [4.3e6*0.7 - 100.,100.,0.,0.,0.,0.,
            4.3e6*0.3 - 100.,100.,0.,0.,0.,0.,
            0.0,0.0,
            0.0,0.0,0.0,0.0]
    f = ODEFunction(twoSEStransmission_with_lower_transmissibility;syms = [:S₁,:E₁,:I₁,:R₁,:W₁,:C₁,:S₂,:E₂,:I₂,:R₂,:W₂,:C₂,:C_sero₁,:C_sero₂,:EW₁,:IW₁,:EW₂,:IW₂])
    return ODEProblem(f,u0,(0,365))
end

"""
function simple_conv(incidence,kern)

Direct implementation of convolution [incidence ⋆ kern](t)
"""
function simple_conv(incidence,kern)
    z = similar(incidence)
    n = length(incidence)
    m = length(kern)
    for t = 1:n
        z[t] = 0.
        for s = 1:(t-1)
            if t-s <= m && t-s >= 1
                z[t] += kern[t-s]*incidence[s]
            end
        end
    end
    return z
end


"""
Subroutine dgpstep

This subroutine computes the gradient projection step

  s = P[x + alpha*w] - x,

where P is the projection on the n-dimensional interval [xl,xu].
"""
function dgpstep(n,x,xl,xu,alpha,w,s)
    # TODO
    # This computation of the gradient projection step avoids
    # rouding errors for the components that are feasible.

    @inbounds for i=1:n
        if x[i] + alpha*w[i] < xl[i]
            s[i] = xl[i] - x[i]
        elseif x[i] + alpha*w[i] > xu[i]
            s[i] = xu[i] - x[i]
        else
            s[i] = alpha*w[i]
        end
    end
    return
end