include("../diagonal_sbp.jl")
include("legacy_functions.jl")

if length(ARGS) != 0
    level = parse(Int,ARGS[1])
    Iterations = parse(Int,ARGS[2])
else
    level = 10
    Iterations = 10000
end

using LinearAlgebra
using SparseArrays
using Plots
using IterativeSolvers
using BenchmarkTools
using MAT

function e(i,n)
    # A = Matrix{Float64}(I,n,n)
    # return A[:,i]
    out = spzeros(n)
    out[i] = 1.0
    return out
end

function eyes(n)
    # return Matrix{Float64}(I,n,n)
    out = spzeros(n,n)
    for i in 1:n
        out[i,i] = 1.0
    end
    return out
end

function u(x,y)
    return sin.(π*x .+ π*y)
end

function Diag(A)
    # Self defined function that is similar to Matlab Diag
    return Diagonal(A[:])
end

h_list_x = [1/2^1, 1/2^2, 1/2^3, 1/2^4, 1/2^5, 1/2^6, 1/2^7, 1/2^8, 1/2^9, 1/2^10, 1/2^11, 1/2^12, 1/2^13, 1/2^14]
h_list_y = [1/2^1, 1/2^2, 1/2^3, 1/2^4, 1/2^5, 1/2^6, 1/2^7, 1/2^8, 1/2^9, 1/2^10, 1/2^11, 1/2^12, 1/2^13, 1/2^14]

function Operators_2d(i, j, h_list_x, h_list_y; p=2)
    hx = h_list_x[i];
    hy = h_list_y[j];

    x = range(0,step=hx,1);
    y = range(0,step=hy,1);
    m_list = 1 ./h_list_x;
    n_list = 1 ./h_list_y;

    # Matrix Size
    N_x = Integer(m_list[i]);
    N_y = Integer(n_list[j]);

    (D1x, HIx, H1x, r1x) = diagonal_sbp_D1(p,N_x,xc=(0,1));
    (D2x, S0x, SNx, HI2x, H2x, r2x) = diagonal_sbp_D2(p,N_x,xc=(0,1));

    (D1y, HIy, H1y, r1y) = diagonal_sbp_D1(p,N_y,xc=(0,1));
    (D2y, S0y, SNy, HI2y, H2y, r2y) = diagonal_sbp_D2(p,N_y,xc=(0,1));

    # BSx = sparse(SNx - S0x);
    # BSy = sparse(SNy - S0y);
    BSx = SNx - S0x
    BSy = SNy - S0y

    # Forming 2d Operators
    # e_1x = sparse(e(1,N_x+1));
    # e_Nx = sparse(e(N_x+1,N_x+1));
    # e_1y = sparse(e(1,N_y+1));
    # e_Ny = sparse(e(N_y+1,N_y+1));
    e_1x = e(1,N_x+1);
    e_Nx = e(N_x+1,N_x+1);
    e_1y = e(1,N_x+1);
    e_Ny = e(N_y+1,N_y+1);

    # I_Nx = sparse(eyes(N_x+1));
    # I_Ny = sparse(eyes(N_y+1));
    I_Nx = eyes(N_x+1);
    I_Ny = eyes(N_y+1);

    e_E = kron(e_Nx,I_Ny);
    e_W = kron(e_1x,I_Ny);
    e_S = kron(I_Nx,e_1y);
    e_N = kron(I_Nx,e_Ny);

    E_E = kron(sparse(Diag(e_Nx)),I_Ny);   # E_E = e_E * e_E'
    E_W = kron(sparse(Diag(e_1x)),I_Ny);
    E_S = kron(I_Nx,sparse(Diag(e_1y)));
    E_N = sparse(kron(I_Nx,sparse(Diag(e_Ny))));

    D1_x = kron(D1x,I_Ny);
    D1_y = kron(I_Nx,D1y);

    D2_x = kron(D2x,I_Ny);
    D2_y = kron(I_Nx,D2y);
    D2 = D2_x + D2_y

    HI_x = kron(HIx,I_Ny);
    HI_y = kron(I_Nx,HIy);

    H_x = kron(H1x,I_Ny);
    H_y = kron(I_Nx,H1y);

    BS_x = kron(BSx,I_Ny);
    BS_y = kron(I_Nx,BSy);

    HI_tilde = kron(HIx,HIx);
    H_tilde = kron(H1x,H1y);

    return (D1_x, D1_y, D2_x, D2_y, D2, HI_x, HI_y, BS_x, BS_y, HI_tilde, H_tilde, I_Nx, I_Ny, e_E, e_W, e_S, e_N, E_E, E_W, E_S, E_N)
end

function Assembling_matrix(level;p=2)
    i = j = level
    hx = h_list_x[i];
    hy = h_list_y[j];

    x = range(0,step=hx,1);
    y = range(0,step=hy,1);
    m_list = 1 ./h_list_x;
    n_list = 1 ./h_list_y;

    N_x = Integer(m_list[i]);
    N_y = Integer(n_list[j]);

    Nx = N_x + 1;
    Ny = N_y + 1;

    (D1_x, D1_y, D2_x, D2_y, D2, HI_x, HI_y, BS_x, BS_y, HI_tilde, H_tilde, I_Nx, I_Ny, e_E, e_W, e_S, e_N, E_E, E_W, E_S, E_N) = Operators_2d(i,j,h_list_x,h_list_y,p=p);
    analy_sol = u(x,y');

    # Penalty Parameters
    tau_E = -13/hx;
    tau_W = -13/hx;
    tau_N = -1;
    tau_S = -1;

    beta = 1;

    # Forming SAT terms

    ## Formulation 1
    SAT_W = tau_W*HI_x*E_W + beta*HI_x*BS_x'*E_W;
    SAT_E = tau_E*HI_x*E_E + beta*HI_x*BS_x'*E_E;

    # SAT_S = tau_S*HI_y*E_S*D1_y
    # SAT_N = tau_N*HI_y*E_N*D1_y

    SAT_S = tau_S*HI_y*E_S*BS_y;
    SAT_N = tau_N*HI_y*E_N*BS_y;

    SAT_W_r = tau_W*HI_x*E_W*e_W + beta*HI_x*BS_x'*E_W*e_W;
    SAT_E_r = tau_E*HI_x*E_E*e_E + beta*HI_x*BS_x'*E_E*e_E;
    SAT_S_r = tau_S*HI_y*E_S*e_S;
    SAT_N_r = tau_N*HI_y*E_N*e_N;

    (alpha1,alpha2,alpha3,alpha4,beta) = (tau_N,tau_S,tau_W,tau_E,beta);

    g_W = sin.(π*y);
    g_E = -sin.(π*y);
    g_S = -π*cos.(π*x);
    # g_N = -π*cos.(π*x)
    g_N = π*cos.(π*x .+ π);

    # Solving with CPU
    A = D2 + SAT_W + SAT_E + SAT_S + SAT_N;

    b = -2π^2*u(x,y')[:] + SAT_W_r*g_W + SAT_E_r*g_E + SAT_S_r*g_S + SAT_N_r*g_N;

    A = -H_tilde*A;
    b = -H_tilde*b;

    return (A,b,H_tilde,Nx,Ny)
end

function prolongation_matrix(N)
    # SBP preserving
    # N = 2^level + 1
    odata = spzeros(2*N-1,N)
    for i in 1:2*N-1
        if i % 2 == 1
            odata[i,div(i+1,2)] = 1
        else
            odata[i,div(i,2)] = 1/2
            odata[i,div(i,2)+1] = 1/2
        end
    end
    return odata
end

function restriction_matrix(N)
    # SBP preserving
    odata = spzeros(div(N+1,2),N)
    odata[1,1] = 1/2
    odata[1,2] = 1/2
    odata[end,end-1] = 1/2
    odata[end,end] = 1/2
    for i in 2:div(N+1,2)-1
        odata[i,2*i-2] = 1/4
        odata[i,2*i-1] = 1/2
        odata[i,2*i] = 1/4
    end
    return odata
end

function restriction_matrix_normal(N)
    # SBP preserving
    odata = spzeros(div(N+1,2),N)
    odata[1,1] = 1/2
    odata[1,2] = 1/4
    odata[end,end-1] = 1/4
    odata[end,end] = 1/2
    for i in 2:div(N+1,2)-1
        odata[i,2*i-2] = 1/4
        odata[i,2*i-1] = 1/2
        odata[i,2*i] = 1/4
    end
    return odata
end

function prolongation_2d(N)
    prolongation_1d = prolongation_matrix(N)
    prolongation_2d = kron(prolongation_1d,prolongation_1d)
    return prolongation_2d
end

function restriction_2d(N)
    restriction_1d = restriction_matrix_normal(N)
    # restriction_1d = restriction_matrix(N)
    restriction_2d = kron(restriction_1d,restriction_1d)
    return restriction_2d
end

function restriction_2d_GPU(N)
    return CUDA.CUSPARSE.CuSparseMatrixCSC(restriction_2d(N))
end

function prolongation_2d_GPU(N)
    return CUDA.CUSPARSE.CuSparseMatrixCSC(prolongation_2d(N))
end

function CG_CPU(A,b,x;maxiter=length(b),abstol=sqrt(eps(real(eltype(b)))),direct_sol=0,H_tilde=0)
    r = b - A * x;
    p = r;
    rsold = r' * r
    # Ap = spzeros(length(b))
    Ap = similar(b);

    num_iter_steps = 0
    norms = [sqrt(rsold)]
    errors = []
    if direct_sol != 0 && H_tilde != 0
        append!(errors,sqrt(direct_sol' * A * direct_sol))
    end
    # @show rsold
    for step = 1:maxiter
        num_iter_steps += 1
        mul!(Ap,A,p);
        alpha = rsold / (p' * Ap)
        x .= x .+ alpha * p;
        r .= r .- alpha * Ap;
        rsnew = r' * r
        append!(norms,sqrt(rsnew))
        if direct_sol != 0 && H_tilde != 0
            error = sqrt((x - direct_sol)' * A * (x - direct_sol))
            append!(errors,error)
        end
        if sqrt(rsnew) < abstol
              break
        end
        p = r + (rsnew / rsold) * p;
        rsold = rsnew
    end
    return (num_iter_steps,norms,errors)
end

function regularCG!(A,b,x,H1,exact;maxiter=length(b)^2,abstol=sqrt(eps(real(eltype(b)))))
    r = b - A * x;
    p = r;
    rsold = r'*r

    num_iter_steps = 0
    norms = [sqrt(r'*H1*r)]

    diff = x-exact
    err = sqrt(diff'*A*diff)

    E = [err]
    for step = 1:maxiter
        Ap = A*p;
        num_iter_steps += 1
        alpha = rsold/(p'*Ap)
        x .= x .+ alpha * p;

        diff = x-exact
        err = sqrt(diff'*A*diff)
        append!(E, err)

        r .= r .- alpha * Ap;
        rsnew = r' * r

        norm_r = sqrt(rsnew)
        append!(norms,norm_r)

        if norm_r < abstol
            break
        end

        beta = rsnew/rsold;
        p = r + beta * p;
        rsold = rsnew
    end

    return E, num_iter_steps, norms
end

function jacobi_brittany!(x,A,b;maxiter=3, ω = 2/3)

    Pinv = Diagonal(1 ./ diag(A))
    P = Diagonal(diag(A))
    Q = A-P

    for j in 1:maxiter
        x[:] = ω * Pinv*(b .- Q*x[:]) + (1 - ω)*x[:]
    end
end

function modified_richardson!(x,A,b;maxiter=3,ω=0.15)
    for _ in 1:maxiter
        x[:] .= x[:] + ω*(b .- A*x[:])
    end
end

function modified_richardson_GPU!(x_GPU,A_GPU,b_GPU;maxiter=3,ω=0.15)
    for _ in 1:maxiter
        x_GPU .= x_GPU .+ ω*(b_GPU .- A_GPU*x_GPU)
    end
end

function Two_level_multigrid(A,b,Nx,Ny,A_2h;nu=3,NUM_V_CYCLES=1,SBPp=2)
    v_values = Dict(1=>zeros(Nx*Ny))
    rhs_values = Dict(1 => b)
    N_values = Dict(1 => Nx)
    N_values[2] = div(Nx+1,2)

    x = zeros(length(b));
    v_values[1] = x

    for cycle_number in 1:NUM_V_CYCLES
        # jacobi_brittany!(v_values[1],A,b;maxiter=nu);
        modified_richardson!(v_values[1],A,b,maxiter=nu)
        r = b - A*v_values[1];
        f = restriction_2d(Nx) * r;
        v_values[2] = A_2h \ f

        # println("Pass first part")
        e_1 = prolongation_2d(N_values[2]) * v_values[2];
        v_values[1] = v_values[1] + e_1;
        # println("After coarse grid correction, norm(A*x-b): $(norm(A*v_values[1]-b))")
        # jacobi_brittany!(v_values[1],A,b;maxiter=nu);
        modified_richardson!(v_values[1],A,b,maxiter=nu)
    end
    return (v_values[1],norm(A * v_values[1] - b))
end

function Two_level_multigrid_GPU(A_GPU,b_GPU,Nx,Ny,A_2h;nu=3,NUM_V_CYCLES=1,SBPp=2)
    v_values = Dict(1=>CuArray(zeros(Nx*Ny)))
    rhs_values = Dict(1 => b_GPU)
    N_values = Dict(1 => Nx)
    N_values[2] = div(Nx+1,2)

    x = CuArray(zeros(length(b_GPU)));
    v_values[1] = x

    for cycle_number in 1:NUM_V_CYCLES
        # jacobi_brittany!(v_values[1],A,b;maxiter=nu);
        modified_richardson_GPU!(v_values[1],A_GPU,b_GPU,maxiter=nu)
        r = b_GPU - A_GPU*v_values[1];
        f = Array(restriction_2d_GPU(Nx) * r);
        v_values[2] = CuArray(A_2h \ f)

        # println("Pass first part")
        e_1 = prolongation_2d_GPU(N_values[2]) * v_values[2];
        v_values[1] = v_values[1] + e_1;
        # println("After coarse grid correction, norm(A*x-b): $(norm(A*v_values[1]-b))")
        # jacobi_brittany!(v_values[1],A,b;maxiter=nu);
        modified_richardson_GPU!(v_values[1],A_GPU,b_GPU,maxiter=nu)
    end
    return (v_values[1],norm(A_GPU * v_values[1] - b_GPU))
end

function precond_matrix(A, b, A_2h; m=3, solver="jacobi")
    #pre and post smoothing
    N = length(b)
    IN = sparse(Matrix(I, N, N))
    P = Diagonal(diag(A))
    Pinv = Diagonal(1 ./ diag(A))
    Q = P-A
    L = A - triu(A)
    U = A - tril(A)

    if solver == "jacobi"
       ω = 2/3
        H = ω*Pinv*Q + (1-ω)*IN
        R = ω*Pinv
        R0 = ω*Pinv
    elseif solver == "ssor"
        ω = 1.4  #this is just a guess. Need to compute ω_optimal (from jacobi method)
        B1 = (P + ω*U)\Matrix(-ω*L + (1-ω)*P)
        B2 = (P + ω*L)\Matrix(-ω*U + (1-ω)*P)
        H = B1*B2
        X = (P+ω*L)\Matrix(IN)

        R = ω*(2-ω)*(P+ω*U)\Matrix(P*X)
        R0 = ω*(2-ω)*(P+ω*U)\Matrix(P*X)
    elseif solver == "richardson"
        ω =ω_richardson
        H = IN - ω*A
        R = ω*IN
        # R = spzeros(N,N)
        R0 = ω*IN
    elseif solver == "richardson_chebyshev" #TODO: FIX ME FOR CHEB
        ω =ω_richardson
        H = IN - ω*A
        R = ω*IN
        R0 = ω*IN
    elseif solver == "chebyshev" #TODO: FIX ME FOR CHEB
        ω =ω_richardson
        H = IN - ω*A
        R = ω*IN
        R0 = ω*IN
    else
    end

    for i = 1:m-1
        R += H^i * R0
    end

    # (A_2h, b_2h, x_2h, H1_2h) = get_operators(SBPp, 2*h);
    A_2h = A_2h
    I_r = standard_restriction_matrix_2D(N)

    I_p = standard_prolongation_matrix_2D(length(b_2h))
    M = H^m * (R + I_p * (A_2h\Matrix(I_r*(IN - A * R)))) + R
    M_alternative = R + R + I_p * (A_2h \ Matrix(I_r * (IN - A*R)))

    M_v0 = I_p * (A_2h \ Matrix(I_r)) # No richardson iteration
    M_v1_no_post = ω*IN + I_p * (A_2h \ Matrix(I_r * (IN - ω*A)))

    M_v1 = (IN - ω*A)*(ω*IN + I_p * (A_2h \ Matrix(I_r * (IN - ω*A)))) + ω*IN # one pre and post richardson iteration

    M_v1_alternative = H * (R0 + I_p * (A_2h \ Matrix(I_r * H))) + R0 # alternative form

    M_test = H^m * R + R + H^m * (I_p * (A_2h \ Matrix(I_r * H^m))) # need to change, a generalized representation of M for richardson iteration

    # return (M, R, H, I_p, A_2h, I_r, IN)
    return (M_test,R,H,I_p,A_2h,I_r,IN)
end

function mg_preconditioned_CG(A,b,x;A_2h = A_2h_lu,maxiter=length(b),abstol=sqrt(eps(real(eltype(b)))),NUM_V_CYCLES=1,nu=3,use_galerkin=true,direct_sol=0,H_tilde=0,SBPp=2)
    Nx = Ny = Int(sqrt(length(b)))
    level = Int(log(2,Nx-1))
    # (A_2h,b_2h,H_tilde_2h,Nx_2h,Ny_2h) = Assembling_matrix(level-1,p=SBPp);
    # A_2h = lu(A_2h)
    r = b - A * x;
    # (M, R, H, I_p, A_2h, I_r, IN) = precond_matrix(A,b;m=nu,solver="jacobi",p=p)
    z = Two_level_multigrid(A,r,Nx,Ny,A_2h;nu=nu,NUM_V_CYCLES=1)[1]
    # z = M*r
    p = z;
    num_iter_steps = 0
    norms = [norm(r)]
    errors = []
    if direct_sol != 0 && H_tilde != 0
        append!(errors,sqrt(direct_sol' * A * direct_sol))
    end

    rzold = r'*z

    for step = 1:maxiter
    # for step = 1:5
        num_iter_steps += 1
        alpha = rzold / (p'*A*p)
        x .= x .+ alpha * p;
        r .= r .- alpha * A*p
        rs = r' * r
        append!(norms,sqrt(rs))
        if direct_sol != 0 && H_tilde != 0
            error = sqrt((x - direct_sol)' * A * (x - direct_sol))
            # @show error
            append!(errors,error)
        end
        if sqrt(rs) < abstol
            break
        end
        z = Two_level_multigrid(A,r,Nx,Ny,A_2h;nu=nu,NUM_V_CYCLES=1)[1]
        # z = M*r
        rznew = r'*z
        beta = rznew/(rzold);
        p = z + beta * p;
        rzold = rznew
    end
    # @show num_iter_steps
    return num_iter_steps, norms, errors
end

function mg_preconditioned_CG_GPU(A_GPU,b_GPU,x_GPU;A_2h = A_2h_lu,maxiter=length(b),abstol=sqrt(eps(real(eltype(b)))),NUM_V_CYCLES=1,nu=3,use_galerkin=true,direct_sol=0,H_tilde=0,SBPp=2)
    Nx = Ny = Int(sqrt(length(b_GPU)))
    level = Int(log(2,Nx-1))
    # (A_2h,b_2h,H_tilde_2h,Nx_2h,Ny_2h) = Assembling_matrix(level-1,p=SBPp);
    # A_2h = lu(A_2h)
    r_GPU = b_GPU - A_GPU * x_GPU;
    # (M, R, H, I_p, A_2h, I_r, IN) = precond_matrix(A,b;m=nu,solver="jacobi",p=p)
    z_GPU = Two_level_multigrid_GPU(A_GPU,r_GPU,Nx,Ny,A_2h;nu=nu,NUM_V_CYCLES=1)[1]
    # z = M*r
    p_GPU = z_GPU;
    num_iter_steps = 0
    norms = [norm(r_GPU)]
    errors = []
    if direct_sol != 0 && H_tilde != 0
        append!(errors,sqrt(direct_sol' * A * direct_sol))
    end

    rzold = r_GPU'*z_GPU

    for step = 1:maxiter
    # for step = 1:5
        num_iter_steps += 1
        alpha = rzold / (p_GPU'*A_GPU*p_GPU)
        x_GPU .= x_GPU .+ alpha * p_GPU;
        r_GPU .= r_GPU .- alpha * (A_GPU*p_GPU)
        rs = r_GPU' * r_GPU
        append!(norms,sqrt(rs))
        # if direct_sol != 0 && H_tilde != 0
        #     error = sqrt((x_GPU - direct_sol)' * A_GPU * (x_GPU - direct_sol))
        #     # @show error
        #     append!(errors,error)
        # end
        if sqrt(rs) < abstol
            break
        end
        z_GPU = Two_level_multigrid_GPU(A_GPU,r_GPU,Nx,Ny,A_2h;nu=nu,NUM_V_CYCLES=1)[1]
        # z = M*r
        rznew = r_GPU'*z_GPU
        beta = rznew/(rzold);
        p_GPU .= z_GPU .+ beta * p_GPU;
        rzold = rznew
    end
    # @show num_iter_steps
    return num_iter_steps, norms, errors
end

function test_preconditioned_CG(;level=level,nu=3,ω=2/3,SBPp=2)
    (A,b,H_tilde,Nx,Ny) = Assembling_matrix(level,p=SBPp);
    (A_2h,b_2h,H_tilde_2h,Nx_2h,Ny_2h) = Assembling_matrix(level-1,p=SBPp);
    A_2h_lu = lu(A_2h)
    direct_sol = A\b
    reltol = sqrt(eps(real(eltype(b))))
    x = zeros(Nx*Ny);
    abstol = norm(A*x-b) * reltol
    h = 1/(Nx-1)
    ω_richardson = 0.15
    p = SBPp
    (M, R, H, I_p, A_2h, I_r, IN) = precond_matrix(A,b,A_2h;m=nu,solver="richardson")

    cond_A_M = cond(M*A)
    x = zeros(Nx*Ny);
    iter_mg_cg, norm_mg_cg, error_mg_cg = mg_preconditioned_CG(A,b,x;A_2h = A_2h_lu, maxiter=length(b),abstol=abstol,NUM_V_CYCLES=1,nu=nu,use_galerkin=true,direct_sol=direct_sol,H_tilde=H_tilde,SBPp=SBPp)
    error_mg_cg_bound_coef = (sqrt(cond_A_M) - 1) / (sqrt(cond_A_M) + 1)
    error_mg_cg_bound = error_mg_cg[1] .* 2 .* error_mg_cg_bound_coef .^ (0:1:length(error_mg_cg)-1)
    scatter(log.(10,error_mg_cg),label="error_mg_cg", markercolor = "darkblue")
    xlabel!("Iterations")
    ylabel!("Log errors, base 10")
    plot!(log.(10,error_mg_cg_bound),label="error_mg_cg_bound",linecolor = "darkblue")

    cond_A = cond(Matrix(A))
    x0 = zeros(Nx*Ny)
    (E_cg, num_iter_steps_cg, norms_cg) = regularCG!(A,b,x0,H_tilde,direct_sol;maxiter=20000,abstol=abstol)

    scatter!(log.(10,E_cg),label="error_cg", markercolor = "darksalmon")
    error_cg_bound_coef = (sqrt(cond_A) - 1) / (sqrt(cond_A) + 1)
    error_cg_bound = E_cg[1] .* 2 .* error_cg_bound_coef .^ (0:1:length(E_cg)-1)
    plot!(log.(10,error_cg_bound),label="error_cg_bound", linecolor = "darksalmon")

    savefig("convergence_richardson.png")
    # my_solver = "jacobi"
    # x0 = zeros(Nx*Ny)
    # (E_mgcg, num_iter_steps_mgcg, norms_mgcg) = MGCG!(A,b,x0,H_tilde,direct_sol,my_solver;smooth_steps = nu,maxiter=20000,abstol=abstol)

    # plot(error_mg_cg,label="error_mg_cg")
    # plot!(error_mg_cg_bound,label="error_mg_cg_bound")
end


using Pkg
Pkg.activate(@__DIR__)
using Soss
using ZigZagBoomerang

include("zigzag.jl")

m = @model x begin
    α ~ Normal()
    β ~ Normal()
    yhat = α .+ β .* x
    y ~ For(eachindex(x)) do j
        Normal(yhat[j], 2.0)
    end
end

x = randn(20);
obs = -0.1 .+ 2x + 1randn(20);
posterior = m(x=x) | (y=obs,)

T = 100.0
trace, final, (num, acc) = @time zigzag(posterior, T)

# trace is a continous object, discretize to obtain samples
ts, xs = ZigZagBoomerang.sep(discretize(trace, 0.1))

xs = xform(posterior).(xs)

using Plots
p = plot(ts, getindex.(xs, 1))
plot!(ts, getindex.(xs, 2), color=:red)
p2 = plot(first.(xs), last.(xs))
p


const obsnames = ["Energy", "Energy^2", "Specific Heat",
                  "|Magnetization|", "Magnetization^2", "Magnetization^4",
                  "Susceptibility", "Connected Susceptibility", "Binder Ratio",
                 ]
const Ns = [2,8]
const J = 1.0
const Ts = [0.3, 1.0, 3.0, 10.0]

"""
Exact calculation for Ising model on `N` site fully connected network
"""
function exact(J, N, Ts)
    j = J/N
    Es = zeros(N+1)
    Ms = zeros(N+1)
    ns = zeros(Int,N+1)
    for Nup in 0:N
        Ndown = N - Nup
        n = binomial(N, Nup)
        push!(ns, n)
        E = j*(Nup*Ndown - 0.5(Nup*(Nup-1)) - 0.5(Ndown*(Ndown-1)))
        push!(Es, E)
        M = (Nup-Ndown)/N
        push!(Ms, M)
    end
    sorted = sortperm(Es, rev=true)
    E0 = Es[sorted[end]]

    mbetas = Ts .\ (-1.0)
    nT = length(Ts)
    res = Dict(n=>zeros(nT) for n in obsnames)
    res["Magnetization"] = zeros(nT)
    Zs = zeros(nT)
    for i in sorted
        E = Es[i] - E0
        zs = exp.(mbetas.*E) * ns[i]
        Zs .+= zs
        E += E0
        res["Energy"] .+= zs.*(E/N)
        res["Energy^2"] .+= zs.*(E/N)^2

        M = Ms[i]
        res["Magnetization"] .+= zs.*M
        res["|Magnetization|"] .+= zs.*abs(M)
        res["Magnetization^2"] .+= zs.*M^2
        res["Magnetization^4"] .+= zs.*M^4
    end
    res["Magnetization"] ./= Zs
    for n in obsnames
        res[n] ./= Zs
    end
    return res
end

betas = 1.0 ./ Ts

for N in Ns
    res = exact(J,N,Ts)
    res["Specific Heat"] = N.*(res["Energy^2"] .- res["Energy"].^2).*betas.^2
    res["Susceptibility"] = N.*(res["Magnetization^2"] .- res["Magnetization"].^2).*betas
    res["Connected Susceptibility"] = N.*(res["Magnetization^2"] .- res["|Magnetization|"].^2).*betas
    res["Binder Ratio"] = res["Magnetization^4"] ./ (res["Magnetization^2"].^2)
    open(@sprintf("J_%.1f__N_%d.dat", J, N), "w") do io
        for i in 1:length(Ts)
            @printf(io, "%.15f", Ts[i])
            for n in obsnames
                @printf(io, " %.15f", res[n][i])
            end
            println(io)
        end
    end
end


using Compat, ClassicalLaminateTheory

forcesAndMoments =
 @compat Dict(
    :Nx => 300,
    :Ny => 100,
    :Nxy => 200,
    :Mx => 5,
    :My => 3,
    :Mxy => 1)

materialProperties =@compat Dict(
  :E1 => 20.0*10^6,
  :E2 => 1.5*10^6,
  :G12 => 1.0*10^6,
  :ν12 => .34,
  :alpha1 => 0.2e-6,
  :alpha2 => 20e-6,
  :beta1 => 0.2e-4,
  :beta2 => 20e-4,
  :thickness => 0.0075)

laminateProperties =@compat Dict(
  :nplies => 5,
  :repeats => 1,
  :symmetric => false,
  :orientation => [0, 45, 90, 45, -45],
  :deltaTemp => -270.0,
  :deltaMoisture => 1
)

createLaminate!(laminateProperties, materialProperties, forcesAndMoments)

fm = [
157.552,
-42.4476,
210.226,
4.84661,
3.15339,
0.846607
]

@assert(round(laminateProperties[:fTotal], 3) == round(fm, 3))
println("Tests for CLT2 passed!")


module TestRandomInitialParameters
    verbose = false

    using Base.Test

    import IntertemporalChoiceHeuristics:
        load,
        fit,
        models,
        initial_parameters,
        training_losses

    import Distributions: Uniform

    import Optim

    inputs, weights = load()

    λ = 0.0

    @testset "Identical optima found from different starting points" begin
        for model in models
            Θ = initial_parameters(model)
            K = length(Θ)
            for loss in training_losses
                res = fit(model, inputs, weights, loss, λ, Θ)

                fₓ = Optim.minimum(res)

                n_points = 10
                costs = Array(Float64, n_points)
                for i in 1:n_points
                    ϵ = rand(Uniform(-0.1, +0.1), K)
                    Θ′ = Θ + ϵ
                    try
                        res = fit(model, inputs, weights, loss, λ, Θ′)
                        fₓ′ = Optim.minimum(res)
                        costs[i] = fₓ′
                        δ = (fₓ′ - fₓ) / fₓ
                        if verbose
                            @printf(
                                "%s\t%s\t%d\t%.4f\t%.4f\t%.4f\n",
                                model,
                                loss,
                                i,
                                fₓ,
                                fₓ′,
                                δ,
                            )
                        end
                    catch
                        costs[i] = fₓ
                    end
                end
                # TODO: Get this to pass reliably.
                # @test all(x -> x < 1e-10, (costs ./ minimum(costs)) - 1.0)
            end
        end
    end
end


# SPDX-License-Identifier: MIT

const Abstract3DTensor{T} = AbstractArray{T,3} where T

#======================================================================
Forward recursion
======================================================================#

function αrecursion!(α::AbstractMatrix{T}, π::AbstractVector{T},
                     Aᵀ::AbstractMatrix{T}, lhs::AbstractMatrix) where T
    N = size(lhs, 2)
    buffer = similar(lhs[:, 1])
    @views elmul!(α[:,1], π, lhs[:,1])
    @views for n in 2:N
        matmul!(buffer, Aᵀ, α[:,n-1])
        elmul!(α[:,n], buffer, lhs[:,n])
    end
    α
end

function αrecursion!(α::Abstract3DTensor{T}, π::AbstractVector{T},
                     Aᵀ::AbstractMatrix{T}, lhs::Abstract3DTensor{T}) where T
    N = size(lhs, 2)
    buffer = similar(α[:,1,:])
    @views elmul!(α[:,1,:], π, lhs[:,1,:])
    @views for n in 2:N
        matmul!(buffer, Aᵀ, α[:,n-1,:])
        elmul!(α[:,n,:], buffer, lhs[:,n,:])
    end
    α
end

#======================================================================
Backward recursion
======================================================================#

function βrecursion!(β::AbstractMatrix{T}, ω::AbstractVector{T},
                     A::AbstractMatrix{T}, lhs::AbstractMatrix{T}) where T
    N = size(lhs, 2)
    buffer = similar(lhs[:, 1])
    β[:, end] = ω
    @views for n in N-1:-1:1
        elmul!(buffer, β[:,n+1], lhs[:,n+1])
        matmul!(β[:,n], A, buffer)
    end
    β
end

function βrecursion!(β::Abstract3DTensor{T}, ω::AbstractVector{T},
                     A::AbstractMatrix{T}, lhs::Abstract3DTensor{T}) where T
    N = size(lhs, 2)
    buffer = fill!(similar(β[:,1,:]), one(T))
    @views elmul!(β[:,end,:], ω, buffer)
    @views for n in N-1:-1:1
        elmul!(buffer, β[:,n+1,:], lhs[:,n+1,:])
        matmul!(β[:,n,:], A, buffer)
    end
    β
end

#======================================================================
Generic forward-backward algorithm
======================================================================#

function αβrecursion(π::AbstractVector{T}, ω::AbstractVector{T},
                     A::AbstractMatrix{T}, Aᵀ::AbstractMatrix,
                     lhs::AbstractMatrix{T}) where T
    S, N = length(π), size(lhs, 2)
    α = similar(lhs, T, S, N)
    β = similar(lhs, T, S, N)

    αrecursion!(α, π, Aᵀ, lhs)
    βrecursion!(β, ω, A, lhs)

    γ = similar(lhs, T, S, N)
    elmul!(γ, α, β)
    sums = sum(γ, dims = 1)
    eldiv!(γ, γ, sums)

    γ, minimum(sums)
end

function αβrecursion(π::AbstractVector{T}, ω::AbstractVector{T},
                     A::AbstractMatrix{T}, Aᵀ::AbstractMatrix,
                     lhs::Abstract3DTensor{T}) where T
    S, N, B = size(lhs)
    α = similar(lhs, T, S, N, B)
    β = similar(lhs, T, S, N, B)

    αrecursion!(α, π, Aᵀ, lhs)
    βrecursion!(β, ω, A, lhs)

    γ = similar(lhs, T, S, N, B)
    elmul!(γ, α, β)
    sums = sum(γ, dims = 1)
    eldiv!(γ, γ, sums)

    γ, dropdims(minimum(sums, dims = (1, 2)), dims = (1, 2))
end

#======================================================================
Specialized algorithms
======================================================================#

_convert(T, ttl::Number) = convert(T, ttl)
_convert(T, ttl::AbstractVector) = copyto!(similar(ttl, T), ttl)

function stateposteriors(in_cfsm, in_lhs::AbstractArray{T}) where T <: Real
    SF = LogSemifield{T}
    cfsm = convert(CompiledFSM{SF}, in_cfsm)
    lhs = copyto!(similar(in_lhs, SF), in_lhs)
    γ, ttl = αβrecursion(cfsm.π, cfsm.ω, cfsm.A, cfsm.Aᵀ, lhs)
    out = copyto!(similar(in_lhs), γ)
    exp.(out), _convert(T, ttl)
end

function bestpath(in_cfsm, in_lhs::AbstractArray{T}) where T <: Real
    SF = TropicalSemifield{T}
    cfsm = convert(CompiledFSM{SF}, in_cfsm)
    lhs = copyto!(similar(in_lhs, SF), in_lhs)
    γ, _ = αβrecursion(cfsm.π, cfsm.ω, cfsm.A, cfsm.Aᵀ, lhs)
    out = copyto!(similar(in_lhs), γ)
    idxs = mapslices(argmax, out, dims = 1)
    dropdims(idxs, dims = 1), dropdims(lhs[idxs], dims = 1)
end


# Linear Discriminant Analysis

mutable struct LinearDiscriminantModel{T} <: DiscriminantModel{T}
    "Standard discriminant model parameters"
    Θ::DiscriminantParameters{T}
    "Whitening transformation"
    W::AbstractMatrix{T}
    "Canonical coordinate projection matrix with whitening"
    C::Union{Nothing,AbstractMatrix{T}}
    "Discriminant function intercept"
    δ₀::T
    function LinearDiscriminantModel{T}() where T
        new{T}(DiscriminantParameters{T}(), Matrix{T}(undef,0,0), nothing, zero(T))
    end
end

function canonical_coordinates!(LDA::LinearDiscriminantModel{T}) where T
    Θ = LDA.Θ

    m, p = check_dims(Θ.M, dims=Θ.dims)

    if p ≤ m-1
        # no dimensionality reduction is possible
        LDA.C = Matrix{T}(I, p, p)
    elseif Θ.dims == 1
        # center M by overall centroid
        M = broadcast!(-, similar(Θ.M, T), Θ.M, transpose(Θ.μ))

        # Σ_between = Mᵀdiag(π)M so need to scale by sqrt π
        broadcast!((πₖ, Mₖⱼ) -> √(πₖ)Mₖⱼ, M, Θ.π, M)
        UDVᵀ = svd!(M*LDA.W, full=false)

        # Extract m-1 components
        LDA.C = LDA.W*transpose(view(UDVᵀ.Vt, 1:m-1, :))
    else
        M = broadcast!(-, similar(Θ.M, T), Θ.M, Θ.μ)

        broadcast!((πₖ, Mⱼₖ) -> √(πₖ)Mⱼₖ, M, transpose(Θ.π), M)
        UDVᵀ = svd!(LDA.W*M, full=false)

        LDA.C = transpose(view(UDVᵀ.U, :, 1:m-1))*LDA.W
    end

    return LDA
end

"""
    fit!(LDA::LinearDiscriminantModel, Y::Vector, X::Matrix; dims::Integer=1, kwargs...)

Fit a linear discriminant model based on data matrix `X` and class indicator matrix `Y`
along dimensions `dims` and overwrite `LDA`.

# Keyword arguments
- `canonical::Bool=false`: computes canonical coordinates and performs dimensionality
  reduction if `true`
- `compute_covariance`: if true, computes `Σ`
- `centroids::Matrix`: specifies the class centroids. If `dims` is `1`, then each row
   represents a class centroid, otherwise each column represents a class centroid. If not
   specified, the centroids are computed from the data
- `priors::Vector`: specifies the class membership prior probabilties. If not specified,
  the class probabilities a computed based on the class counts
- `gamma::Real=0`: regularization parameter γ ∈ [0,1] shrinks the within-class covariance
  matrix towards the identity matrix scaled by the average eigenvalue of the covariance
  matrix
"""
function _fit!(LDA::LinearDiscriminantModel{T},
               y::Vector{<:Integer},
               X::Matrix{T},
               dims::Integer=1,
               canonical::Bool=false,
               compute_covariance::Bool=false,
               centroids::Union{Nothing,AbstractMatrix}=nothing,
               priors::Union{Nothing,AbstractVector}=nothing,
               gamma::Union{Nothing,Real}=nothing) where T
    Θ = LDA.Θ

    set_dimensions!(Θ, y, X, dims)
    set_gamma!(Θ, gamma)
    set_statistics!(Θ, y, X, centroids)
    set_priors!(Θ, priors)

    # Compute degress of freedom
    df = size(X, dims) - Θ.m

    if dims == 1
        # Overall centroid is prior-weighted average of class centroids
        Θ.μ = vec(transpose(Θ.π)*Θ.M)

        # Center the data matrix with respect to classes to compute whitening matrix
        X .-= view(Θ.M, y, :)

        # Compute covariance matrix if requested
        if compute_covariance
            Θ.Σ = lmul!(one(T)/df, transpose(X)*X)
        end
    else
        Θ.μ = Θ.M*Θ.π

        X .-= view(Θ.M, :, y)

        if compute_covariance
            Θ.Σ = lmul!(one(T)/df, X*transpose(X))
        end
    end

    # Use cholesky whitening if gamma is not specifed, otherwise svd whitening
    if Θ.Σ !== nothing
        LDA.W, LDA.δ₀ = whiten_cov_chol!(copy(Θ.Σ), Θ.γ, dims=dims)
    elseif Θ.γ === nothing
        LDA.W, LDA.δ₀ = whiten_data_chol!(X, dims=dims, df=df)
    else
        LDA.W, LDA.δ₀ = whiten_data_svd!(X, Θ.γ, dims=dims, df=df)
    end

    # Perform canonical discriminant analysis if applicable
    if canonical
        canonical_coordinates!(LDA)
    else
        LDA.C = nothing
    end

    Θ.fit = true

    return LDA
end

function lda(X::Matrix{T},
             y::Vector{<:Integer};
             dims::Integer=1,
             canonical::Bool=false,
             compute_covariance::Bool=false,
             centroids::Union{Nothing,AbstractMatrix}=nothing,
             priors::Union{Nothing,AbstractVector}=nothing,
             gamma::Union{Nothing,Real}=nothing) where T
    U = T == Float32 || T == Float64 ? T : Float64
    return _fit!(
        LinearDiscriminantModel{U}(),
        y,
        U == T ? copy(X) : convert(Matrix{U}, X),
        dims,
        canonical,
        compute_covariance,
        centroids,
        priors,
        gamma
    )
end

function discriminants!(Δ::Matrix{T}, LDA::LinearDiscriminantModel{T}, X::Matrix{T}) where T
    dims = LDA.Θ.dims

    M = LDA.Θ.M

    n, p = check_dims(X, dims=dims)
    m, p₂ = check_dims(M, dims=dims)
    n₂, m₂ = check_dims(Δ, dims=dims)

    alt_dims = dims == 1 ? 2 : 1

    p == p₂ || throw(DimensionMismatch("predictor count along dimension $(alt_dims) of " *
                                       "X must match dimension $(dims) of M (got $(p) " *
                                       "and $(p₂))"))
    n == n₂ || throw(DimensionMismatch("observation count along dimension $(dims) of X " *
                                       "must match dimension $(dims) of Δ (got $(n) " *
                                       "and $(n₂))"))
    m == m₂ || throw(DimensionMismatch("class count along dimension $(alt_dims) of Δ " *
                                       "must match dimension $(dims) of M (got $(m) " *
                                       "and $(m₂))"))

    is_row = dims == 1

    π = LDA.Θ.π
    W = LDA.C === nothing ? LDA.W : LDA.C

    X̃, M̃ = is_row ? (X*W, M*W) : (W*X, W*M)

    Z = similar(X̃)
    for k = 1:m
        if is_row
            Δₖ = view(Δ, :, k:k)
            μₖ = view(M̃, k:k, :)
        else
            Δₖ = view(Δ, k:k, :)
            μₖ = view(M̃, :, k:k)
        end

        broadcast!(-, Z, X̃, μₖ)
        sum!(abs2, Δₖ, Z)
        broadcast!((d², logπ) -> logπ - d²/2, Δₖ, Δₖ, log(π[k]))
    end

    return Δ
end


compute(n::Int)::Int = (n - 5) ÷ 7 * 3 + 2


using Test
import LessOLS: Process, make_sample, sampler_normal, ols, desc, show, yhat, equation

# # use profiler
# # https://thirld.com/blog/2015/05/30/julia-profiling-cheat-sheet/

gen = sampler_normal(a=0, b=20, β_0=5, β=[1, 2.5], sd_e=0.5)
sam = gen(100)
lm = ols(sam, intercept=true)
show(lm)

@assert equation([-1/3, -2/7, 1/29],true,4) == "Y = -0.3333 - 0.2857*X1 + 0.0345*X2"


# Tracing charged particle in Tokamak.
# This example shows how to trace protons in a stationary magnetic field that
# corresponds to a Tokamak.
#
# Hongyang Zhou, hyzhou@umich.edu

using TestParticle
using TestParticle: getB_Tokamak
using OrdinaryDiffEq
using StaticArrays
using PyPlot
using Statistics: mean
using Printf

## Obtain field

# Magnetic bottle parameters in SI units
const ICoil = 80. # current in the coil
const N = 15000 # number of windings
const IPlasma = 1e6 # current in the plasma
const a = 1.5 # radius of each coil
const b = 0.8 # radius of central region

function getB(xu)
   SVector{3}(getB_tokamak(xu[1], xu[2], xu[3], a, b, ICoil*N, IPlasma))
end

function getE(xu)
   SA[0.0, 0.0, 0.0]
end

## Initialize particles
m = TestParticle.mᵢ
q = TestParticle.qᵢ
c = TestParticle.c

# initial velocity, [m/s]
v₀ = [-0.1, -0.15, 0.0] .* c
# initial position, [m]
r₀ = [2.3, 0.0, 0.0]
stateinit = [r₀..., v₀...]

param = prepare(getE, getB; species=Proton)
tspan = (0.0, 1e-6)

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

@printf "Speed = %6.4f %s\n" √(v₀[1]^2+v₀[2]^2+v₀[3]^2)/c*100 "% speed of light"
@printf "Energy = %6.4f MeV\n" (1/√(1-(v₀[1]/c)^2-(v₀[2]/c)^2-(v₀[3]/c)^2)-1)*m*c^2/abs(q)/1e6

# Default Tsit5() alone does not work in this case! You need to set a maximum
# timestep to maintain stability, or choose a different algorithm as well.
# The sample figure in the gallery is generated with AB3() and dt=2e-11.
sol = solve(prob, Tsit5(); dt=2e-11, save_idxs=[1,2,3])

## Visualization

using3D()
fig = plt.figure(figsize=(10,6))
ax = fig.gca(projection="3d")

ax.plot(sol[1,:], sol[2,:], sol[3,:], label="proton")

# Plot coils
θ = range(0, 2π, length=201)
y = a * cos.(θ)
z = a * sin.(θ)
for i in 0:17
   ϕ = i*π/9
   ax.plot(y*sin(ϕ) .+ (a+b)*sin(ϕ), y*cos(ϕ) .+ (a+b)*cos(ϕ), z, color="k")
end

# Plot Tokamak
u = range(0, 2π, length=100)
v = range(0, 2π, length=100)

U = [y for _ in u, y in v]
V = [x for x in u, _ in v]

X = @. (a + b + (a - 0.05)*cos(U)) * cos(V)
Y = @. (a + b + (a - 0.05)*cos(U)) * sin(V)
Z = @. (a - 0.05) * sin(U)

ax.plot_surface(X, Y, Z, alpha=0.1)

ax.legend()
title("Particle trajectory in Tokamak")
xlabel("x [m]")
ylabel("y [m]")
zlabel("z [m]")

xMin, xMax = -4, 4
yMin, yMax = -4, 4
zMin, zMax = -4, 4

xlim(xMin, xMax)
ylim(yMin, yMax)
zlim(zMin, zMax)

TestParticle.set_axes_equal(ax)

#savefig("test.png", bbox_inches="tight")


#=
Copyright 2022 INSIGNEO Institute for in silico Medicine

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
=#

"""
    setInletBC(t :: Float64, dt :: Float64, v :: Vessel,  h  :: Heart)

Flow or pressure value is set from the inlet file stored in the `Heart` structure.
"""
function setInletBC(t :: Float64, dt :: Float64, v :: Vessel)
    h = v.heart

  	if h.inlet_type == "Q"
		v.Q[1] = inputFromData(t, h)
  	else
		v.P[1] = inputFromData(t, h)
  	end

	inletCompatibility(dt, v, h)
end

"""
    inputFromData(t :: Float64, h :: Heart)

Inlet flow waveform is stored in `h` data structure inside `input_data` matrix.
`input_data` has two columns and as many rows as time steps.
"""
function inputFromData(t :: Float64, h :: Heart)
	idt = h.input_data[:,1]
	idq = h.input_data[:,2]

	t_hat = div(t,h.cardiac_T)

	t -= t_hat*h.cardiac_T

	idx = 0
	@inbounds for i in 1:length(idt)
		if ((t >= idt[i]) && (t <= idt[i+1]))
			idx=i
			break
		end
	end

	@inbounds qu = idq[idx] + (t - idt[idx]) * (idq[idx+1] - idq[idx]) /
       (idt[idx+1] - idt[idx])

	return qu
end

"""
    inletCompatibility(dt :: Float64, v :: Vessel, h :: Heart)

Compute compatibility relations at inlet node.
"""
function inletCompatibility(dt :: Float64, v :: Vessel, h :: Heart)
	W11, W21 = riemannInvariants(1, v)
	W12, W22 = riemannInvariants(2, v)

	@fastmath @inbounds W11 += (W12 - W11)*(v.c[1] - v.u[1])*dt*v.invDx
	@fastmath @inbounds W21 = 2.0*v.Q[1]/v.A[1] - W11

	v.u[1], v.c[1] = inverseRiemannInvariants(W11, W21)

	if h.inlet_type == "Q"
		@fastmath @inbounds v.A[1] = v.Q[1]/v.u[1]
		v.P[1] = pressure(v.A[1], v.A0[1], v.beta[1], v.Pext)
	else
		v.A[1] = areaFromPressure(v.P[1], v.A0[1], v.beta[1], v.Pext)
		@fastmath @inbounds v.Q[1] = v.u[1]*v.A[1]
	end

end

"""
    riemannInvariants(i :: Int, v :: Vessel)

Calculate Riemann invariants at the node `i` from `u` and `c`.
"""
function riemannInvariants(i :: Int, v :: Vessel)
  @fastmath @inbounds W1 = v.u[i] - 4.0*v.c[i]
  @fastmath @inbounds W2 = v.u[i] + 4.0*v.c[i]

  return W1, W2
end

"""
    inverseRiemannInvariants(W1 :: Float64, W2 :: Float64)

Calculate `u` and `c` given `W1` and `W2`
"""
function inverseRiemannInvariants(W1 :: Float64, W2 :: Float64)
  @fastmath u = 0.5*(W1 + W2)
  @fastmath c = (W2 - W1)*0.125

  return u, c
end

"""
    areaFromPressure(P :: Float64, A0 :: Float64, beta :: Float64, Pext :: Float64)

Inverse constitutive equation. This is used only when a pressure inlet-time-function is
imposed (not recommended).
"""
function areaFromPressure(P :: Float64, A0 :: Float64, beta :: Float64, Pext :: Float64)
   return A0 * ((P-Pext)/beta + 1.0)*((P-Pext)/beta + 1.0)
end

"""
    setOutletBC(dt :: Float64, v :: Vessel)

Outlet boundary condition is applied to the last node of the involved vessel. Two
boundary conditions are available. Either a reflection coefficient is provided or a
three-element windkessel is coupled.
"""
function setOutletBC(dt :: Float64, v :: Vessel)
    if v.outlet == "reflection"
        v.P[end] = 2.0*v.P[end-1] - v.P[end-2]
		outletCompatibility(dt, v)
    elseif v.outlet == "wk3"
        threeElementWindkessel(dt, v)
	end
end

"""
    function outletCompatibility(dt :: Float64, v :: Vessel)

Outlet compatibility relations compute all the quantities not directly assigned by
the outlet boundary condition.
"""
function outletCompatibility(dt :: Float64, v :: Vessel)
	W1M1, W2M1 = riemannInvariants(v.M-1, v)
	W1M, W2M   = riemannInvariants(v.M, v)

	W2M += (W2M1 - W2M)*(v.u[end] + v.c[end])*dt/v.dx
	W1M = v.W1M0 - v.Rt * (W2M - v.W2M0)

	v.u[end], v.c[end] = inverseRiemannInvariants(W1M, W2M)
	v.Q[end] = v.A[end]*v.u[end]
end

"""
    function threeElementWindkessel(dt :: Float64, v :: Vessel)

The three element windkessel simulates the perfusion of downstream vessels.
This 0D model is coupled by `wk3` function to 1D model terminal branches via the
solution of a Riemann problem at the 0D/1D interface.
"""
function threeElementWindkessel(dt :: Float64, v :: Vessel)
	Pout = 0.0

	Al = v.A[end]
	ul = v.u[end]

	v.Pc += dt/v.Cc * (Al*ul - (v.Pc - Pout)/v.R2)

	As = Al

	ssAl = sqrt(sqrt(Al))
	sgamma =  2*sqrt(6*v.gamma[end])
	sA0 = sqrt(v.A0[end])
	bA0 = v.beta[end]/sA0

	fun(As) = As*v.R1*(ul + sgamma * (ssAl - sqrt(sqrt(As)))) -
		(v.Pext + bA0*(sqrt(As) - sA0)) + v.Pc

	dfun(As) = v.R1*(ul + sgamma * (ssAl - 1.25*sqrt(sqrt(As)))) - bA0*0.5/sqrt(As)

    try
        As = newtonSolver(fun, dfun, As)
    catch e
        vlab = v.label
        println("\nNewton solver doesn't converge at $vlab outlet!")
        throw(e)
    end

	us = (pressure(As, v.A0[end], v.beta[end], v.Pext) - Pout)/(As*v.R1)

	v.A[end] = As
	v.u[end] = us
end

"""
    function newtonSolver(f, df, x0 :: Float64)

Solve windkessel equation by means of Newton method.
"""
function newtonSolver(f, df, x0 :: Float64)
	xn = x0 - f(x0)/df(x0)
	if abs(xn-x0)<= 1e-5
		return xn
	else
		newtonSolver(f, df, xn)
	end
end

"""
    updateGhostCells(v :: Vessel)

Ghost cells are updated by assuming both ends as transmissive boundaries.
"""
function updateGhostCells(v :: Vessel)
	v.U00A = v.A[1]
	v.U00Q = v.Q[1]
	v.U01A = v.A[2]
	v.U01Q = v.Q[2]

	v.UM1A = v.A[v.M]
	v.UM1Q = v.Q[v.M]
	v.UM2A = v.A[v.M-1]
	v.UM2Q = v.Q[v.M-1]
end

function updateGhostCells(vessels :: Array{Vessel,1})
	for vessel in vessels
		updateGhostCells(vessel)
	end
end


@testset "CARLIN algorithm: logistic model" begin
    r = -0.5
    K = 0.8
    F1 = hcat(r)
    F2 = hcat(-r/K)
    X0 = 0.47 .. 0.53
    prob = @ivp(CanonicalQuadraticForm(F1, F2), x(0) ∈ X0)

    sol = solve(prob, T=10.0, alg=CARLIN(N=3, bloat=true))
    @test dim(sol) == 1
    @test ρ([1.0], sol(10.0)) < 0.2

    sol = solve(prob, T=10.0, alg=CARLIN(N=3, bloat=false))
    @test dim(sol) == 1
    @test ρ([1.0], sol(10.0)) < 0.01

    # Use the compressed Kronecker form
    sol = solve(prob, T=10.0, alg=CARLIN(N=3, bloat=false, compress=true))
    @test dim(sol) == 1
    @test ρ([1.0], sol(10.0)) < 0.01
end


######################################################
# Created by Alex Dowling (alexdowling.net)
# while at the University of Wisconsin-Madison
# ####################################################

# Define problem
include("../DegeneracyHunter.jl")
using JuMP

# Define problem
function lp1()

	m = Model()

	@variable(m, 0 <= x[1:3] <= 5)
	@variable(m, 0 <= y <= 0)
	@constraint(m, x[1] + x[2] >= 1)
	@constraint(m, x[1] + x[2] + x[3] == 1)
	@constraint(m, x[2] - 2*x[3] <= 1)
	@constraint(m, x[1] + x[3] >= 1)
	@constraint(m, x[1] + x[2] + x[3] == 1) # Redundant constraint - makes problem degenerate
	@objective(m, Min, sum{x[i],i=1:3})

	return m
end

function initialize!(m::Model)
	x = getvariable(m,:x)
	setvalue(x[1], 1.0)
	setvalue(x[2], 5.0)
	setvalue(x[3], -1.0)

	# Intentionally do not initialize y

	return nothing
end

# Create and initialize model
m = lp1()
initialize!(m)

# Check initialization
DegeneracyHunter.printVariableDiagnostics(m)

# Print equations with absolute residuals greater than 0.001
DegeneracyHunter.printInfeasibleEquations(m, 0.001)

# Check for degeneracy constraints
DegeneracyHunter.degeneracyHunter(m, includeBounds=true)

# Solve model
tic()
status = solve(m)
tm = toq()

# Print inactive constraints at solution
DegeneracyHunter.printInactiveEquations(m)

# Package and print problem size information
ps = DegeneracyHunter.assembleProblemStats(m,status,tm)
DegeneracyHunter.printProblemStats(ps)


using Lehmann
using Printf
using Gaston
β = 40.0
rtol = 1e-8
eta = 1e-2
N = 128

dlr = DLRGrid(β = β, Euv = 1.0, isFermi = true, rtol = rtol, symmetry = :none)
τgrid = collect(LinRange(0.0, β, N))
println(τgrid)
G = Sample.SemiCircle(dlr, :τ, τgrid)
G += rand(length(G)) / 2.0 * eta
coeff = tau2dlr(dlr, G, τgrid)
coeff_sumrule = tau2dlr(dlr, G, τgrid, sumrule = π / 2)
println("sum rule: ", abs(sum(coeff) - π / 2), " versus ", abs(sum(coeff_sumrule) - π / 2))

dlr10 = DLRGrid(β = β, Euv = 100.0, isFermi = true, rtol = rtol, symmetry = :none)

Gtrue = Sample.SemiCircle(dlr, :τ, dlr10.τ)
Gfit = real.(dlr2tau(dlr, coeff, dlr10.τ))
Gfit_sumrule = real.(dlr2tau(dlr, coeff_sumrule, dlr10.τ))

@printf("%15s%30s%30s%30s\n", "tau", "true", "no sumrule", "with sumrule")
for i in 1:dlr10.size
    @printf("%15.6f%30.15f%30.15e%30.15e\n", dlr10.τ[i], Gtrue[i], abs(Gfit[i] - Gtrue[i]), abs(Gfit_sumrule[i] - Gtrue[i]))
end

set(term = "qt")
# p = plot(dlr10.τ, Gtrue, label = "original", axis = "semilogy")
xrange = (0.0, 40)
p = plot(dlr10.τ, Gfit - Gtrue, plotstyle = :linespoints, leg = Symbol("DLR no sum rule"), Axes(xrange = xrange))
plot!(dlr10.τ, Gfit_sumrule - Gtrue, plotstyle = :linespoints, leg = Symbol("DLR with sum rule"))
display(p)
readline()


using Printf
using Statistics
using LinearAlgebra
using MDBM
using PyPlot
pygui(true);
using Interpolations
using Plots

dim=2 #DoF of the system (after Cauchy transcription)
mult=3 #Ns=dim*mult
n=40 #timestep number
#choosing numerical simulation method
#RK4: 4th order Runge-Kutta
method="RK4"
i=20 #number of iteration
iT=8

#parameters of the delayed Mathieu equation
tau=2*pi
omega=1
kappa=0.1
delta=1.5
epsilon=1
b=0.25
T=2*pi/omega

#timestep
dt=tau/(n-1)

#matrix coefficients
function A(t)
 [-kappa -(delta+epsilon*cos(omega*t)); 1 0]
end
B=[0 b; 0 0]

#checking the relation of tau and T
function checkT(tau,T)
    if tau <= T
        println("OK")
    else
        println("Error")
    end
end

checkT(tau,T)
k=Int16(floor(T/tau)) #k factor

#initial function ()
vect1=rand(n,mult*dim)*(-10)+ones(n,mult*dim)*5

#interpolate
function it(A)
    matrdim=size(A,2)-1
    scaleitp=real(A[:,1])
    knots=(scaleitp,)
    ARe=real(A)
    AIm=imag(A)

    imRe=[interpolate(knots, ARe[:,2],Gridded(Linear()))]
    for j = 3:matrdim+1
        imRetemp=[interpolate(knots, real(ARe[:,j]), Gridded(Linear()))]
        imRe=vcat(imRe,imRetemp)
    end
    imIm=[interpolate(knots, AIm[:,2],Gridded(Linear()))]
    for j = 3:matrdim+1
        imImtemp=[interpolate(knots, real(AIm[:,j]), Gridded(Linear()))]
        imIm=vcat(imIm,imImtemp)
    end
    return(hcat(imRe,imIm))
end

#substitution in interpolating function
function sub(it,t)
    subdim=size(it,1)
    out=zeros(ComplexF64,subdim)
    for j = 1:subdim
        out[j]=it[j,1](t)+it[j,2](t)*im
    end
    return(out)
end
collect(-tau:dt:iT*tau)
function simulation(init)
        solreturn=zeros(ComplexF64,1,1+mult*dim) #matrix for time integration solution (test)
        eval=zeros(ComplexF64,mult*dim,1) #eigenvalue matrix

        sol00=init #S0 inital random array
        #iteration loop
        for g = 1:i
        sol=zeros(iT*n-iT,mult*dim) #empty solution vector
        tvect0=collect(-tau:dt:iT*tau+dt/2) #discretized time vector
        sol0=hcat(tvect0,vcat(sol00,sol)) #constructing solution array
            #time integration
            for m = 2:dim:(dim*mult-(dim-2)) #choosing "one physical system"
            sol0m=hcat(sol0[:,1],sol0[:,m:m+dim-1])
                #time integration steps
                for j = 0:iT*n-iT-1
                    sol0mit=it(sol0m)

                    Y1=sub(sol0mit,j*dt)
                    Y2=Y1+(dt/2)*(A(j*dt)*Y1+B*sub(sol0mit,-tau+j*dt))
                    Y3=Y1+(dt/2)*(A((j+0.5)*dt)*Y2+B*sub(sol0mit,-tau+(j+0.5)*dt))
                    Y4=Y1+dt*(A((j+0.5)*dt)*Y3+B*sub(sol0mit,-tau+(j+0.5)*dt))
                    Y=Y1+(dt/6)*((A(j*dt)*Y1+B*sub(sol0mit,-tau+j*dt))+2*(A((j+0.5)*dt)*Y2+B*sub(sol0mit,-tau+(j+0.5)*dt))+2*(A((j+0.5)*dt)*Y3+B*sub(sol0mit,-tau+(j+0.5)*dt))+(A((j+1)*dt)*Y4+B*sub(sol0mit,-tau+(j+1)*dt)))

                    sol0m[n+j+1,2:2+(dim-1)]=transpose(Y) #filling solution array of the physical system
                end
                sol0[:,m:m+dim-1]=sol0m[:,2:2+(dim-1)] #filling iterational solution array
            end
            solreturn=vcat(solreturn,sol0) #filling global solution array (test)
            S=sol0[(iT-1)*n-((iT-1)-1):(iT)*n-(iT-1),2:mult*dim+1] #taking array Sj (n x Ns)
            V=sol0[iT*n-(iT-1):(iT+1)*n-iT,2:mult*dim+1] #taking array Vj (n x Ns)
            #H=pinv(S)*V #calculating matrix Hj (Ns x Ns)
            H=pinv(S)*V
            Am=eigvecs(H) #calculation eigenvectors of Hj, creating matrix Am (Ns x Ns)
            sol00=zeros(ComplexF64,n,mult*dim) #calculating matrix Sj+1 (n x Ns)
            for o = 1:n
                sol00[o,:]=Am*V[o,:]
            end
            sol00=sol00
            evm=eigvals(H) #obtaining dominant eigenvalues in the iteration step
            eval=hcat(eval,evm) #filling up the eigenvalue array
        end
        return(eval[:,2:i+1]) #returning the global eigenvalue array
        #return(solreturn)
end

aa=simulation(vect1)

aa[6000:6100,:]

function evplot(aaa,k)
U=aaa[:,k]
eigRe1=[real(U[n,1]) for n in 1:size(U,1)]
eigIm1=[imag(U[n,1]) for n in 1:size(U,1)]
plot1=Plots.plot(eigRe1,eigIm1,seriestype=:scatter)
Re_circle=[cos(n) for n in 0:0.05:2*pi]
Im_circle=[sin(n) for n in 0:0.05:2*pi]
return(Plots.plot!(plot1,Re_circle,Im_circle))
end

evplot(aa,1)
evplot(aa,2)
evplot(aa,3)
evplot(aa,4)
evplot(aa,5)
evplot(aa,6)
evplot(aa,7)
evplot(aa,8)
evplot(aa,9)
evplot(aa,10)
evplot(aa,12)
evplot(aa,11)
evplot(aa,13)
evplot(aa,14)
evplot(aa,15)
evplot(aa,16)
evplot(aa,17)
evplot(aa,18)
evplot(aa,19)
evplot(aa,20)
evplot(aa,21)
evplot(aa,22)
evplot(aa,23)
evplot(aa,24)
evplot(aa,25)
evplot(aa,26)
evplot(aa,27)
evplot(aa,28)
evplot(aa,29)
evplot(aa,30)
evplot(aa,31)
evplot(aa,32)

solt1=real(aa)
v1plot=[solt1[j,2] for j in 1:size(solt1,1)]
v2plot=[solt1[j,3] for j in 1:size(solt1,1)]
v3plot=[solt1[j,4] for j in 1:size(solt1,1)]
v4plot=[solt1[j,5] for j in 1:size(solt1,1)]
v5plot=[solt1[j,6] for j in 1:size(solt1,1)]
v6plot=[solt1[j,7] for j in 1:size(solt1,1)]
Plots.plot(v1plot)
Plots.plot(v2plot)
Plots.plot(v3plot)
Plots.plot(v4plot)
Plots.plot(v5plot)
Plots.plot(v6plot)

function avarr(A)
    dim1=size(A,1)
    dim2=size(A,2)
    sum=0
    for i=1:dim1
            for j=1:dim2
                sum=sum+norm(A[i,j])
            end
    end
    return(sum/(dim1*dim2))
end

avarr([1 2; 3 4])
(1/2.5)*[1 2; 3 4]

norm(-1)
avarr([1 2+4*im; 3 4])


function newtondist!(yout::Vector{Float64},iters::Vector{Int64},Vs::Float64,vs::Float64,ts::Float64,
    V::Float64,C::Float64,Q::Float64,V0::Float64,beta::Float64,c0::Float64,A0::Float64,W1end::Float64,rho::Float64,
    f::Function,J::Function,maxiter::Int16,epsJ::Float64,epsN::Float64,maxval::Float64,h::Float64)

    # initial guess
    x0 = zeros(2);
    xx = zeros(2);
    x0[2] = V;
    Pest = 1/C.*(x0[2] .- V0);
    Aest = (Pest./beta .+ sqrt.(A0)).^2;
    x0[1] = W1end .+ 8.*(c0 .- sqrt.(0.5.*beta./rho).*Aest.^0.25);

    # non-dimensionalize
    x0[1] /= vs;
    x0[2] /= Vs;

    # setup for iterations
    xx .= x0;
    x = zeros(length(x0));
    N = 1;

    # max step size for line searches
    stpmax = 100*max(sqrt(norm(xx)),length(x0));
    # println(stpmax)

    # Newton iterations
    while N <= maxiter
        # determine Jacobian, check invertibility
        JJ = J(xx,Vs,vs,ts,C,W1end,c0,rho,beta,h);
        D = diagm(maximum!(zeros(length(xx)),abs.(JJ)).^-1);
        # println(JJ)
        # println(D)
        # println(D*JJ)
        # println(cond(D*JJ))
        # println(det(D*JJ))
        if abs(det(D*JJ)) < epsJ
            println(JJ)
            println(D*JJ)
            print(det(D*JJ))
            error("Distal Newton Jacobian is singular.");
        end
        # compute gradient of line search objective function
        fvec = D*f(xx,V,Q,Vs,vs,ts,V0,rho,beta,C,W1end,c0,A0,h);
        # println(f(xx,system,n,state))
        # println(fvec)
        fold = 0.5*dot(fvec,fvec);
        g = (D*JJ)'*fvec;
        # compute newton step
        # println(inv(D*JJ))
        s = -inv(D*JJ)*fvec;
        # line search to update state vector
        fn, xn, check = CVModule.linedist(xx,fold,g,s,stpmax,f,J,maxiter,V,Q,Vs,vs,ts,V0,rho,beta,C,W1end,c0,A0,h);
        # println(xn[1]*vs)
        # check if sufficiently close to root
        JJ = J(xn,Vs,vs,ts,C,W1end,c0,rho,beta,h);
        D = diagm(maximum!(zeros(length(xx)),abs.(JJ)).^-1);
        fvec = D*f(xn,V,Q,Vs,vs,ts,V0,rho,beta,C,W1end,c0,A0,h);
        # println(fvec)
        if norm(fvec) <= epsN
            x = xn;
            x[1] = x[1]*vs;
            x[2] = x[2]*Vs;
            # println(x)
            # println(system.branches.W2root)
            # println(inv(D*JJ))
            # println(fvec)
            iters[1] += N;
            break
        end
        if check
            test = 0.
            den = max(fn,0.5*length(xn));
            for i = 1:length(xn)
                temp = abs(g[i])*max(abs(xn[i]),1.)/den;
                if temp > test
                    test = temp;
                end
            end
            if test < 1e-6
                check = true;
                println("Warning: gradient of objective function close to zero.")
            else
                check = false;
            end
        end
        # check for divergence
        if norm(fvec,Inf) >= maxval
            iters[1] += N;
            println(xn)
            println(fvec)
            error("Newton iteration diverged.");
        end
        N+=1;
        xx = xn;
        if N == maxiter
            println(JJ)
            # println(D)
            # println(D*JJ)
            println(xn)
            println(fvec)
            println(norm(fvec))
            error("Newton iteration failed to converge.");
        end
    end

    # update based on converged solution
    yout[1] = x[2];
    yout[2] = 1./C.*(x[2] .- V0);
    yout[3] = (2.*rho/beta).^2.*((W1end.-x[1])./8 .+ c0).^4;
    yout[4] = yout[3].*0.5.*(W1end .+ x[1]);
end


x  = -5.0:0.1:5.0

f(x) = 1.0 ./ (1 .+ exp.(-x))

println("Start Plot")
using Plots
plot(x,f(x))

println("save Plot")
savefig("sigmoid.png")


using BenchmarkTools
using LinearAlgebra

SUITE = BenchmarkGroup()
for n in [100, 400, 800]
    A = randn(n, n)
    B = randn(n, n)
    C = zero(A)
    SUITE["n=$n"] = @benchmarkable mul!($C, $A, $B)
end


# Aodwt.jl

using MIRT: Aodwt
using MIRTjim: jim

using LinearMapsAA: LinearMapAM, LinearMapAO
using Test: @test

"""
    Aodwt_show( ; dims::Dims=(64,32), level::Int=3)
show scales
"""
function Aodwt_show( ; dims::Dims = (64, 32), level::Int=3)
	W, scales, mfun = Aodwt(dims, level=level)
	jim(scales)
end

for op in (true, false)
	W,_,_ = Aodwt((8,16) ; level=2, operator=op)
	@test Matrix(W)' == Matrix(W') # check adjoint
	op && isinteractive() && (@show W.wt)
	op && isinteractive() && (@show propertynames(W.wt))
	@test W.level == 2
	@test W.wt.name == "haar"
	@test W isa (op ? LinearMapAO : LinearMapAM)
end

Aodwt_show()


include("intcode.jl")

# Part 1: run the intcode program on all points in the 50x50 starting grid and count 1's
prog = open(f -> parse.(Int64, split(readline(f), ",")), "day19_input.txt")
result = map(x -> run_intcode!(copy(prog), [x[1], x[2]])[1], CartesianIndices((0:49, 0:49)))
println("Number of on pixels: $(sum(result))")

# Part 2: find top left coord of part of beam covering 100x100 area. Estimate first and last y as fn of x
approx_start_factor = 38/50
approx_end_factor = 49/50
test_point(prog, x, y) = run_intcode!(copy(prog), [x, y])[1]
function find_exact_start_end(prog, x)
    ys = Int32(floor(x * approx_start_factor)) - 1
    while test_point(prog, x, ys) == 0 ys += 1 end
    ye = Int32(ceil(x * approx_end_factor)) + 1
    while test_point(prog, x, ye) == 0 ye -= 1 end
    return ys, ye
end
# Look for the first coordinate which can fit a 100x100 square in
for x in 700:1200
    ys1, ye1 = find_exact_start_end(prog, x)
    ys2, ye2 = find_exact_start_end(prog, x + 99)
    if ye1 - ys2 >= 99
        println("First coord for square: ($x, $ys2). Code: $(x * 10000 + ys2)")
        break
    end
end


function test_convolution_layer(backend::Backend, n_group, filter_w, filter_h, pad_w, pad_h, stride_w, stride_h, n_input, freeze, T, eps)
  println("-- Testing Convolution(frozen=$freeze) on $(typeof(backend)){$T} filter=$((filter_w,filter_h))...")
  println("    > Setup")
  input_w = 16
  input_h = 10
  input_chann = 6
  input_num = 24
  n_filter = 12

  input_dims = (input_w, input_h, input_chann, input_num)
  filter_dims = (filter_w, filter_h, round(Int, input_chann/n_group), n_filter)
  bias_dims = (1, 1, n_filter, 1)

  layer = ConvolutionLayer(name="conv", kernel=(filter_w, filter_h), stride=(stride_w, stride_h),
      pad=(pad_w,pad_h), n_filter=n_filter, n_group=n_group,
      tops=Array(Symbol,n_input), bottoms=Array(Symbol,n_input))

  # convolution layer requires each input blob to be the same shape
  input = [rand(T, input_dims) for i = 1:n_input]
  inputs = Blob[make_blob(backend, x) for x in input]
  data_diffs = Blob[make_blob(backend, x) for x in input]

  state = setup(backend, layer, inputs, data_diffs)

  println("    > Forward")
  filter = rand(T, filter_dims)
  copy!(state.filter, filter)
  bias = rand(T, bias_dims)
  copy!(state.bias, bias)

  forward(backend, state, inputs)
  for i = 1:n_input
    expected_output = convolution_forward(state, filter, bias, input[i])
    @test size(expected_output) == size(state.blobs[i])
    got_output = similar(expected_output)
    copy!(got_output, state.blobs[i])
    @test all(-eps .< expected_output - got_output .< eps)
  end

  println("    > Backward")
  top_diff = [rand(T, size(state.blobs[1])) for i = 1:n_input]
  for i = 1:n_input
    copy!(state.blobs_diff[i], top_diff[i])
  end

  if freeze
    freeze!(state)
  end
  backward(backend, state, inputs, data_diffs)

  grad_filter_exp = zeros(T, filter_dims)
  grad_bias_exp = zeros(T, bias_dims)
  for i = 1:n_input
    gradients_expected = convolution_backward(state, filter, bias, input[i], top_diff[i])
    gradients_got = similar(gradients_expected[3])
    copy!(gradients_got, data_diffs[i])

    @test all(-eps .< gradients_got - gradients_expected[3] .< eps)

    grad_filter_exp += gradients_expected[1]
    grad_bias_exp += gradients_expected[2]
  end
  grad_filter_got = similar(grad_filter_exp)
  grad_bias_got = similar(grad_bias_exp)
  copy!(grad_filter_got, state.∇filter)
  copy!(grad_bias_got, state.∇bias)

  #is_grad_filter_match = all(abs(grad_filter_exp - grad_filter_got) .< eps)
  #is_grad_bias_match   = all(abs(grad_bias_exp - grad_bias_got) .< eps)

  is_grad_filter_match = all(map((x,y)->isapprox(x,y; atol=eps), grad_filter_exp, grad_filter_got))
  is_grad_bias_match   = all(map((x,y)->isapprox(x,y; atol=eps), grad_bias_exp, grad_bias_got))

  if freeze
    # when frozen, the gradients are not computed, so the got value
    # should be random un-initialized values, so should not match (with
    # very high probability)
    @test !is_grad_bias_match
    @test !is_grad_filter_match
  else
    @show maximum(abs(grad_filter_exp - grad_filter_got))
    @show eps
    @test is_grad_filter_match
    @test is_grad_bias_match
  end

  shutdown(backend, state)
end

# naive implementation of convolution forward, used to check the correctness
function convolution_forward(state, filter::Array, bias::Array, input::Array)
  width, height, channel, num = size(input)
  n_group = state.layer.n_group
  o_g = round(Int, state.layer.n_filter / n_group)
  k_g = round(Int, channel / n_group)

  output = Array(eltype(input), size(state.blobs[1]))
  output[:] = 0

  for n = 1:num
    for g = 1:n_group
      for o = 1:o_g
        for k = 1:k_g
          for y = 1:state.height_out
            for x = 1:state.width_out
              for p = 1:state.layer.kernel[2]
                for q = 1:state.layer.kernel[1]
                  in_y = (y-1) * state.layer.stride[2] - state.layer.pad[2] + p
                  in_x = (x-1) * state.layer.stride[1] - state.layer.pad[1] + q
                  if (in_y >= 1 && in_y <= height && in_x >= 1 && in_x <= width)
                    output[x, y, (g-1)*o_g+o, n] += input[in_x, in_y, (g-1)*k_g+k, n] *
                        filter[q, p, k, (g-1)*o_g+o]
                  end
                end
              end
            end
          end
        end
      end
    end
  end

  # add bias
  for n = 1:num
    for o = 1:state.layer.n_filter
      for y = 1:state.height_out
        for x = 1:state.width_out
          output[x, y, o, n] += bias[o]
        end
      end
    end
  end

  return output
end

# naive implementation of convolution backward, used to check correctness
function convolution_backward(state, filter::Array, bias::Array, input::Array, top_diff::Array)
  ∇filter = zeros(eltype(filter), size(filter))
  ∇bias   = zeros(eltype(bias), size(bias))
  ∇input  = zeros(eltype(input), size(input))

  width, height, channels, num = size(input)
  n_group = state.layer.n_group
  o_g = round(Int, state.layer.n_filter / n_group)
  k_g = round(Int, channels / n_group)

  # ∇bias
  for n = 1:num
    for o = 1:state.layer.n_filter
      for y = 1:state.height_out
        for x = 1:state.width_out
          ∇bias[o] += top_diff[x, y, o, n]
        end
      end
    end
  end

  # ∇filter and ∇input
  for n = 1:num
    for g = 1:n_group
      for o = 1:o_g
        for k = 1:k_g
          for y = 1:state.height_out
            for x = 1:state.width_out
              for p = 1:state.layer.kernel[2]
                for q = 1:state.layer.kernel[1]
                  in_y = (y-1) * state.layer.stride[2] - state.layer.pad[2] + p
                  in_x = (x-1) * state.layer.stride[1] - state.layer.pad[1] + q
                  if (in_y >= 1 && in_y <= height && in_x >= 1 && in_x <= width)
                    ∇filter[q,p,k,(g-1)*o_g+o] += top_diff[x,y,(g-1)*o_g+o,n] * input[in_x,in_y,(g-1)*k_g+k,n]
                    ∇input[in_x,in_y,(g-1)*k_g+k,n] += top_diff[x,y,(g-1)*o_g+o,n] * filter[q,p,k,(g-1)*o_g+o]
                  end
                end
              end
            end
          end
        end
      end
    end
  end

  return (∇filter, ∇bias, ∇input)
end

function test_convolution_layer(backend::Backend, n_group, filter_w, filter_h, pad_w, pad_h, stride_w, stride_h, n_input, T, eps)
  test_convolution_layer(backend, n_group, filter_w, filter_h, pad_w, pad_h, stride_w, stride_h, n_input, true, T, eps)
  test_convolution_layer(backend, n_group, filter_w, filter_h, pad_w, pad_h, stride_w, stride_h, n_input, false, T, eps)
end
function test_convolution_layer(backend::Backend, n_input, T, eps)
  test_convolution_layer(backend, 2, 3, 4, 2, 2, 1, 2, n_input, T, eps)
  #XXX: temporarily disable this test for GPU of a bug in cuDNN v2 RC2
  # see https://devtalk.nvidia.com/default/topic/806940/gpu-accelerated-libraries/cudnn-1x1-filter-convolutions
  #if !test_gpu
    test_convolution_layer(backend, 1, 1, 1, 0, 0, 1, 1, n_input, T, eps)
  #end
end

function test_convolution_layer(backend::Backend)
  test_convolution_layer(backend, 3, Float64, 1e-5)
  test_convolution_layer(backend, 3, Float32, 1e-2) # Float32 is sooo inaccurate?
end
if test_cpu
  test_convolution_layer(backend_cpu)
end
if test_gpu
  test_convolution_layer(backend_gpu)
end


using Test
using Miletus

using Miletus.TermStructure
using Miletus.DayCounts

yc = ConstantYieldCurve(Actual365(), 0.01, :Continuous, :NoFrequency, Date("2016-11-01"))
core = CoreModel(40.00, yc, 0.0)
eucall = EuropeanCall(Date("2016-11-25"), SingleStock(), 41.00)
euput = EuropeanPut(Date("2016-11-25"), SingleStock(), 41.00)

σ = ivol(core, eucall, 1.5)
m = fit(GeomBMModel, core, eucall, 1.5)

@test m.volatility == σ
@test value(m, eucall) ≈ 1.5

coref = CoreForwardModel([Date("2016-11-25")], [42.0], yc)

σf = ivol(coref, eucall, 1.5)
mf = fit(GeomBMModel, coref, eucall, 1.5)

@test mf.volatility == σf
@test value(mf, eucall) ≈ 1.5

σ = ivol(core, euput, 1.5)
m = fit(GeomBMModel, core, euput, 1.5)

@test m.volatility == σ
@test value(m, euput) ≈ 1.5

let
    d1 = Dates.today()
    d2 = d1 + Day(365)
    c = EuropeanCall(d2, SingleStock(), 147.32)
    core = CoreModel(d1, 157.32, 0.01, 0.01)
    gbm = GeomBMModel(core, 0.1)
    v = ivol(core, c, 12, TianModel)
    gbm2 = GeomBMModel(core, v)
    @test isapprox(value(gbm2, c), 12, rtol = 1e-2)
    v = ivol(core, c, 12, CRRModel)
    gbm2 = GeomBMModel(core, v)
    @test isapprox(value(gbm2, c), 12, rtol = 1e-2)
    v = ivol(core, c, 12, JRModel)
    gbm2 = GeomBMModel(core, v)
    @test isapprox(value(gbm2, c), 12, rtol = 1e-2)
end


using Base.Test
using ThreePhotons

#Also tested in test_structure.jl

function test_Dlms(l, m, mb)
    l1 = l
    l2 = l
    m1 = m
    m2 = m
    m1b = mb
    m2b = mb

    sum = 0
    da = 0.1
    for theta = 0:da:pi
      for phi = 0:da:2*pi
        for gamma = 0:da:2*pi
          sum += sin(theta)* Dlms(l1, m1, m1b, theta, phi, gamma)* conj(Dlms(l2, m2, m2b, theta, phi, gamma)) * da^3
        end
      end
    end

    @test real(sum) ≈ (8 * pi ^ 2) / (2l + 1) atol=0.02
end

test_Dlms(4,2,2)


using Statistics

export compute_hypercoef

"""
    compute_hypercoef( model )

compute hyperviscosity coefficient
"""
function compute_hypercoef(model)

    # compute gradient of w using 'gradient' matlab function

    # dxUx,dyUx = gradient(w(:,:,1)',grid.x_ref,grid.y_ref)
    # dxUy,dyUy = gradient(w(:,:,2)',grid.x_ref,grid.y_ref)

    nx, ny = model.grid.nx, model.grid.ny

    update_velocities!(model)

    dxUx = irfft(1im .* model.grid.kx .* model.û_x, nx)
    dyUx = irfft(1im .* model.grid.ky .* model.û_x, nx)
    dxUy = irfft(1im .* model.grid.kx .* model.û_y, nx)
    dyUy = irfft(1im .* model.grid.ky .* model.û_y, nx)

    s = (dxUx .+ dyUy) ./ 2
    d = (dyUx .+ dxUy) ./ 2

    @show lambda_rms = sqrt(mean(s .^ 2 .+ d .^ 2))

    dx, dy = model.grid.dx, model.grid.dy

    # Hyperviscosity coefficient
    coef = 40 * lambda_rms * ((dx+dy)/2π)^model.hv_order

    return coef

end


module Gates
"""
Minimal definition of circuit components provided for testing convenience.
QXZoo contains more comprehensive gate set and features
"""

h() = [[1., 1.] [1., -1.]]./sqrt(2)
x() = [[0., 1.] [1., 0.]]
z() = [[1., 0.] [0., -1.]]
I() = [[1., 0.] [0., 1.]]
cx() = [[1., 0., 0., 0.] [0., 1., 0., 0.] [0., 0., 0., 1.] [0., 0., 1., 0.]]
cz() = [[1., 0., 0., 0.] [0., 1., 0., 0.] [0., 0., 1., 0.] [0., 0., 0., -1.]]
swap() = [[1., 0., 0., 0.] [0., 0., 1., 0.] [0., 1., 0., 0.] [0., 0., 0., 1.]]

end


using BenchmarkTools

using ArnoldiMethod: partialschur, SM, LM
using LinearMaps
using LinearAlgebra
using SparseArrays
using DelimitedFiles

mymatrix(n) = spdiagm(-1 => fill(-1.0, n-1), 0 => fill(2.0, n), 1 => fill(-1.001, n-1))

struct ShiftAndInvert{TA,TB,TT}
    A_lu::TA
    B::TB
    temp::TT
end

function (M::ShiftAndInvert)(y,x)
    mul!(M.temp, M.B, x)
    # @show typeof(y), typeof(M.A_lu), typeof(M.temp)
    ldiv!(y, M.A_lu, M.temp)
end

function construct_linear_map(A,B)
    a = ShiftAndInvert(factorize(A),B,Vector{eltype(A)}(undef, size(A,1)))
    LinearMap(a, size(A,1), ismutating=true)
end

# For inverting in non-generealized case
function construct_linear_map(A)
    a = factorize(A)
    function f(y,x)
        ldiv!(y, a, x)
    end
    LinearMap{eltype(A)}(f, size(A,1))
end

function bencharnoldi()
    A = mymatrix(6000)
    A = construct_linear_map(mymatrix(6000))
    target = LM()

    @benchmark partialschur(B, mindim=11, maxdim=22, nev=10, tol=1e-10, restarts=100000, which=tar) setup = (B = $A; tar = $target)

end


file_dir = Pkg.dir("D4M")*"/examples/2Apps/1EntityAnalysis/Entity.jld";

using JLD

E = load(file_dir)["E"]

E = logical(E);

A = sqIn(E);

d = diag(Adj(A))

A = putAdj(A,Adj(A)-diagm(d));

####

i,j,v = findnz(Adj(A));

An = putAdj(A, sparse(i,j,v ./ min(d[i],d[j])));

####

x = "LOCATION/new york,";

p = StartsWith("PERSON/,");

print( (A[p,x] > 4) & (An[p,x] > 0.3));

####

p0 = "PERSON/john kennedy,";

p1 = Row(A[p,p0] + A[p0,p]);

using PyPlot

spy(A[p1,p1])

p2 = Row( A[p1,p1] - (A[p,p0]+ A[p0,p]));

A[p2,p2] > 1


using Pkg
pkg"activate ."
pkg"add JuMP"
pkg"add Cbc"
pkg"add Juniper"
pkg"add Ipopt"

using JuMP, Cbc, Juniper
# Escolhendo o modelo
model = Model(Cbc.Optimizer)
# Definindo os parâmetros que, por sinal, coloquei aleatórios, pois não sei as entradas
m = 4
n = 4
c = rand(1:4, n,m)
a = rand(0:0, m)
b = rand(1000:1000,n)

# Implementando realmente o modelo
@variable(model,x[ i = 1:n, g = 1:m], Bin) # Variáveis
@variable(model,y[i = 1:n,j = 1:n,g = 1:m], Bin) #Variáveis

#Func. objetivo
@objective(model, Max, sum(c[i,j] * y[i,j,g] for i = 1:n-1, j = i+1:n, g = 1:m))

#Restrições
@constraint(model, [i = 1:n], sum(x[i,g] for g = 1:m) == 1)
@constraint(model, [i = 1:n, g = 1:m], sum(x[i,g] for g = 1:m) ≤ b[g])
@constraint(model, [i = 1:n, g = 1:m], a[g] <= sum(x[i,g] for g = 1:m))
@constraint(model,[i = 1:n-1, j = i+1:n, g = 1:m;1 ≤ i < j ≤ n], y[i,j,g] <= x[i,g])
@constraint(model, [i = 1:n-1, j = i+1:n, g = 1:m;1 ≤ i < j ≤ n], y[i,j,g] <= x[j,g])
@constraint(model, [i = 1:n-1, j = i+1:n, g = 1:m;1 ≤ i < j ≤ n], y[i,j,g] ≥ x[i,g] + x[j,g] - 1)

optimize!(model) # Otimizando o modelo

println("Valor da função objetivo = ", objective_value(model))
value.(x)


module PortfolioOpt

using Distributions
using JuMP
using LinearAlgebra
using LinearAlgebra: dot
import Reexport

include("formulations.jl")
include("mean_variance.jl")
include("mean_variance_robust.jl")
include("mean_variance_dro.jl")
include("sample_based.jl")
include("sample_based_robust.jl")
include("sample_based_stochastic.jl")
include("simple_rules.jl")
include("utils.jl")
include("forecasts.jl")
include("test_utils/testutils.jl")

export AbstractPortfolioFormulation,
    AbstractMeanVariance,
    AbstractSampleBased,
    base_model,
    compute_solution,
    MeanVariance,
    mixed_signals_predict_return,
    portfolio_return!,
    portfolio_variance!,
    po_max_conditional_expectation_limit_predicted_return,
    po_max_predicted_return_limit_conditional_expectation,
    po_max_predicted_return_limit_return,
    po_max_return_limit_variance,
    po_max_utility_return,
    po_min_variance_limit_return,
    predicted_portfolio_return!,
    predicted_portfolio_variance!,
    readjust_volumes!,
    RobustBenTal,
    RobustBertsimas,
    RobustDelague,
    RobustBetina,
    SampleBased,
    # end-to-end
    mean_variance_noRf_analytical,
    max_sharpe,
    equal_weights

    Reexport.@reexport using JuMP

end


using JuMP, Cbc, NPZ

function Solve_func(year, tol)

	println("Solving weights for $year ...\n\n")

	array = npzread(string(year, "_input.npz"))

	A1 = array["A1"]
	A2 = array["A2"]
	b = array["b"]

	model = Model(Cbc.Optimizer)
	set_optimizer_attribute(model, "logLevel", 1)
	N = size(A1)[2]

	@variable(model, r[1:N] >= 0)
	@variable(model, s[1:N] >= 0)

	@objective(model, Min, sum(r[i] + s[i] for i in 1:N))

	# bound on top by tolerance
	@constraint(model, [i in 1:N], r[i] + s[i] <= tol)

	# Ax = b
	@constraint(model, [i in 1:length(b)], sum(A1[i,j] * r[j] + A2[i,j] * s[j]
		                          for j in 1:N) == b[i])

	optimize!(model)
	termination_status(model)

	r_vec = value.(r)
	s_vec = value.(s)

	npzwrite(string(year, "_output.npz"), Dict("r" => r_vec, "s" => s_vec))

	println("\n")

end

year_list = [x for x in 2012:2030]
tol_list = [0.40, 0.38, 0.35, 0.33, 0.30, 0.37, 0.38,
			0.38, 0.39, 0.39, 0.38, 0.40, 0.39, 0.41,
			0.41, 0.42, 0.42, 0.42, 0.42]

# Run solver function for all years and tolerances (in order)
for i in zip(year_list, tol_list)
	Solve_func(i[1], i[2])
end


module Fenkar

using StaticArrays

export fenkar!, fenkar

# define fenkar system
function H(x;k=100.0)
	return 0.5*(1.0 + tanh(k*x))
end
function fenkar!(dx, x, p, t) # in-place, no SA, fast

	# parameters
	@views tsi,tv1m,tv2m,tvp,twm,twp,td,to,tr,xk,uc,uv,ucsi = p

	# fenkar dynamics
	dx[1] = -((x[1]/to)*H(uc-x[1]) + H(x[1]-uc)/tr - (x[2]/td)*H(x[1]-uc)*(1.0-x[1])*(x[1]-uc) - (x[3]/tsi)*H(x[1]-ucsi;k=xk))
	dx[2] = H(uc-x[1])*(1.0-x[2])/(tv1m*H(uv-x[1]) + tv2m*H(x[1]-uv)) - H(x[1]-uc)*x[2]/tvp
	dx[3] = H(uc-x[1])*(1.0-x[3])/twm - H(x[1]-uc)*x[3]/twp

	return nothing
end
function fenkar(x,p,t)		# out-of-place, SA, faster

	# parameters
	@views tsi,tv1m,tv2m,tvp,twm,twp,td,to,tr,xk,uc,uv,ucsi = p

	# fenkar dynamics
	du = -((x[1]/to)*H(uc-x[1]) + H(x[1]-uc)/tr - (x[2]/td)*H(x[1]-uc)*(1.0-x[1])*(x[1]-uc) - (x[3]/tsi)*H(x[1]-ucsi;k=xk))
	dv = H(uc-x[1])*(1.0-x[2])/(tv1m*H(uv-x[1]) + tv2m*H(x[1]-uv)) - H(x[1]-uc)*x[2]/tvp
	dw = H(uc-x[1])*(1.0-x[3])/twm - H(x[1]-uc)*x[3]/twp

	return SA[du,dv,dw]
end

# fenkar system parameters in p[1:13] (BR; from https://doi.org/10.1063/1.166311)
#         tsi,  tv1m,   tv2m,  tvp,  twm,    twp,   td,   to,   tr,   xk,  uc,     uv,  ucsi
#   p = [29.0,  19.6, 1250.0, 3.33, 41.0,  870.0, 0.25, 12.5, 33.3, 10.0, 0.13,  0.04,  0.85]
   p = SA[ 22.0, 333.0,   40.0, 10.0, 65.0, 1000.0, 0.12, 12.5, 25.0, 10.0, 0.13, 0.025,  0.85]
   u0= SA[0.2, 0.5, 0.8]
   tspan = (0.0,1000.0)
end


################################################################################
#                      TEST: Kinematics using Quaternion
################################################################################

for i = 1:samples
    # Create a random quaternion.
    qba0 = Quaternion(randn(), randn(), randn(), randn())
    qba0 = qba0/norm(qba0)

    # Create a random velocity vector.
    wba_a = randn(3)

    # Propagate the initial DCM using the sampled velocity vector.
    Δ   = 0.0001
    qba = copy(qba0)
    for k = 1:1000
        dqba = dquat(qba,vect(qba\wba_a*qba))

        qba  = qba + dqba*Δ
    end

    # In the end, the vector aligned with w must not change.
    v1 = vect(qba0\wba_a*qba0)
    v2 = vect( qba\wba_a*qba)

    @test norm(v2-v1) < 1e-10
end


using Memoize
using GenericSchur
using JLD

"""
    greyencode(n::Integer) = n ⊻ (n >> 1)
returns grey encoding of given integer
"""
greyencode(n::Integer) = n ⊻ (n >> 1)

"""
    inverseM(N)
finds inverse of M matrix of size 2^(N-1)×2^(N-1) with elements M_{i,j} equal to
inner product of binary rep of i (classical) and j (grey code), i,j ∈
[1:2^(N-1)] utilizes algorithm from arxiv:quant-ph/0404089
"""
@memoize function inverseM(N)
    m = 2^(N - 1)
    M = zeros((m, m))
    for i in 1:m
        for j in 1:m
            M[i,j] = (-1)^count_ones((i - 1) .& greyencode(j - 1))
        end
    end
    return inv(M)
end

"""
    stateprep(u, N)

creates Array of CircuitGate objects preparing state `u` over `N` qubits.. uses
algorithm from arxiv:quant-ph/0410066
"""
function stateprep(u, N, n=1)
    cg = CircuitGate[]
    k = 2
    grey = zeros(Int, 2^(N - n))
    for i in 0:N - 2
        grey[2^i:2^(i + 1):end] .= i + n + 1
    end
    grey[end] = N
    θ = real.(atan.(-u[2:2:end] ./ u[1:2:end]) .* 2)
    θ[isnan.(θ)] .= -π
    M = inverseM(N - n + 1)
    θ = M * θ
    for i in 1:2^(N - n)
        push!(cg, circuit_gate((n,), RyGate(θ[i])))
        push!(cg, circuit_gate((n), X, grey[i]))
    end
    return cg
end

"""
    unitary2circuit(m::Matrix{Complex64,2}) where {N}

compiles quantum circuit from given unitary matrix. Implements algorithms from
ff10.1016/j.cpc.2019.107001f, 10.1103/PhysRevA.69.032315, arxiv.org:1003.5760,
arxiv:quant-ph/0404089

# Examples
```@example; setup = :(using RandomMatrices)
M = Stewart(ComplexF64, 4);
cgc = unitary2circuit(M)
```
"""
function unitary2circuit(m::AbstractMatrix{ComplexF64}, N::Union{Int, Nothing}=nothing, wires=nothing)
    m * m' ≈ I || error("Only unitary matrices can be compiled into a quantum circuit")
    s = size(m)

    if isnothing(N)
        N = intlog2(s[1])
    end

    if Diagonal(diag(m)) ≈ m
        return compilediag(diag(m), N, wires)
    end

    if size(m)[1] == 2
        return compile1qubit(m, wires)[1]
    end

    if size(m)[1] == 4
        return compile2qubit(m, N, wires)
    end

    cgc = CircuitGate[]
    QR, τ = qr_unblocked(m)
    cgc = compilediag(diag(QR), N, wires)

    Dg = fill(1.0 + 0.0im, (2^N))
    Dg[1] = -1
    Dgm = compilediag(Dg, N)

    for i in 2^(N) - 1:-1:1
        prepd = CircuitGate[]
        u = zeros(ComplexF64, (2^N))
        u[i + 1:end] = QR[i + 1:end, i]
        u[i] = 1
        u = u ./ norm(u)
        angles = exp.(im .* angle.(u))
        u = u ./ angles
        D = compilediag(angles, N)

        for j in 1:N - 1
            cg = stateprep(u[1:2^(j - 1):end], N, j)
            append!(prepd, cg)
            u = apply(u, cg)
        end
        θ = real(atan(-u[2^(N - 1) + 1] ./ u[1]) .* 2)
        if !isnan(θ)
            cg = circuit_gate((N,), RyGate(θ))
        else
            cg = circuit_gate((N,), RyGate(-π))
        end
        push!(prepd, cg)

        prep = reverse(adjoint.(prepd))
        append!(prep, D)
        prepend!(prepd, adjoint.(D))

        append!(cgc, prepd)
        append!(cgc, Dgm)
        append!(cgc, prep)
    end
    return cgc
end

"""
    qr_unblocked(m::Matrix{Complex64,2})

performs serial qr decomposition on matrix `m`
"""
function qr_unblocked!(m::AbstractMatrix{ComplexF64}, n=1::Integer)
    N = size(m)[1]
    τ = ComplexF64[]
    for i in n:N - 1
        b = norm(m[i,i])
        a = angle(m[i,i])
        push!(τ, 1 + b)
        m[i,i] = -exp(im * a)
        m[i + 1:N, i] =  m[i + 1:N, i] ./ (exp(im * a) * (1 + b))
        m[i + 1:N,i + 1:N] -= (m[i + 1:N, i] * transpose(m[i,i + 1:N]))
    end
    return τ
end

"""
    qr_unblocked(m::Matrix{Complex64,2})

performs serial qr decomposition on matrix `m`
"""
function qr_unblocked(m::AbstractMatrix{ComplexF64}, n=1::Integer)
    m = deepcopy(m)
    N = size(m)[1]
    τ = ComplexF64[]
    for i in n:N - 1
        b = norm(m[i,i])
        a = angle(m[i,i])
        push!(τ, 1 + b)
        m[i,i] = -exp(im * a)
        m[i + 1:N, i] =  m[i + 1:N, i] ./ (exp(im * a) * (1 + b))
        m[i + 1:N,i + 1:N] -= (m[i + 1:N, i] * transpose(m[i,i + 1:N]))
    end
    return m, τ
end

"""
    compile1qubit(m::AbstractMatrix{ComplexF64}, N, wires=nothing)

compiles a U(2) matrix into quantum gates
"""
function compile1qubit(m::AbstractMatrix{ComplexF64}, wires=nothing)
    if isnothing(wires)
        wires = [1]
    end
    cg = CircuitGate[]
    phase = sqrt(m[1,1]*m[2,2] - m[2,1]*m[1,2])
    m  = m / phase
    b = norm.(m)
    if (norm(b[1,1]) + norm(b[2,2])) < 1e-5
        θ2 = -imag(log(m[2,1]/m[1,2]))
        θ1 = 0
        ϕ = (real(m[1,2]/exp(-θ1/2*im)) < 0 ? 3π/2 : π/2) * 2
    elseif (norm(b[1,2]) + norm(b[2,1])) < 1e-5
        θ2 = -imag(log(m[1,1]/m[2,2]))
        θ1 = 0
        ϕ = (real(m[1,1]/exp(-θ1/2*im)) < 0 ? π : 0) * 2
    else
        ϕ = real(acos(sqrt(m[1,1] * m[2,2])) * 2)
        θ1 = imag(log(2im * m[2,1] * m[2,2] / sin(ϕ)))
        θ2 = imag(-log(2im * m[1,1] * m[2,1] / sin(ϕ)))
    end

    if !(norm(θ2) < 1e-5)
        append!(cg, circuit_gate.((wires[1],), [RzGate(θ2), RxGate(ϕ)]))
    else
        append!(cg, circuit_gate.((wires[1],), [RxGate(ϕ)]))
    end

    if !(norm(θ1) < 1e-5)
        append!(cg, circuit_gate.((wires[1],), [RzGate(θ1)]))
    end
    return cg, phase
end

iso_permutations = [
                    [[1,1], [2,2], [3,3], [4,4]],
                    [[2,1], [1,2], [4,3], [3,4]],
                    [[3,1], [4,2], [1,3], [2,4]],
                    [[4,1], [3,2], [2,3], [1,4]],
                    [[2,1], [1,2], [4,3], [3,4]],
                    [[1,1], [2,2], [3,3], [4,4]],
                    [[4,1], [3,2], [2,3], [1,4]],
                    [[3,1], [4,2], [1,3], [2,4]],
                    [[3,1], [4,2], [1,3], [2,4]],
                    [[4,1], [3,2], [2,3], [1,4]],
                    [[1,1], [2,2], [3,3], [4,4]],
                    [[2,1], [1,2], [4,3], [3,4]],
                    [[4,1], [3,2], [2,3], [1,4]],
                    [[3,1], [4,2], [1,3], [2,4]],
                    [[2,1], [1,2], [4,3], [3,4]],
                    [[1,1], [2,2], [3,3], [4,4]],
                    ]

iso_signs = [
            [ 1, 1, 1, 1],
            [ 1,-1, 1,-1],
            [ 1,-1,-1, 1],
            [ 1, 1,-1,-1],
            [ 1,-1,-1, 1],
            [-1,-1, 1, 1],
            [-1,-1,-1,-1],
            [ 1,-1, 1,-1],
            [ 1, 1,-1,-1],
            [ 1,-1,-1, 1],
            [-1, 1,-1, 1],
            [-1,-1,-1,-1],
            [ 1,-1, 1,-1],
            [-1,-1,-1,-1],
            [ 1, 1,-1,-1],
            [-1, 1, 1,-1]
            ]

function _get_permutation(m::AbstractMatrix{Float64}, ::Type{Val{I}}, ::Type{Val{J}}) where {I,J}
    val = 0
    for i in 1:4
        val += m[iso_permutations[4*(J-1)+I][i]...] * iso_signs[4*(J-1)+I][i]
    end
    val/4
end

function _get_col!(m::AbstractMatrix{Float64}, col::Vector{Float64}, col_id::Int)
    for i in 1:4
        col[i] = _get_permutation(m, Val{i}, Val{col_id})
    end
end

function _get_row!(m::AbstractMatrix{Float64}, row::Vector{Float64}, row_id::Int)
    for i in 1:4
        row[i] = _get_permutation(m, Val{row_id}, Val{i})
    end
end

"""
    decomposeSO4(m::AbstractMatrix{ComplexF64})

decomposes a SO(4) matrix into 2 SU(2) matrices, such that `m`∈SO(4), A,B ∈
SU(2) via isoclinic decomposition. m ≊ A ⊗ B under the magic basis homomorphism.
"""
function decomposeSO4(m::AbstractMatrix{Float64})
    size(m)[1] == size(m)[2] || error("decomposeSO4 only works on square matrices")
    size(m)[1] == 4 || error("decomposeSO4 only works on 4x4 matrices")
    det(m) ≈ 1 || error("matrix `m` is not in SO(4), determinant is not 1")

    col = zeros(Float64, 4)
    row = zeros(Float64, 4)

    col_id = 0
    row_id = 0

    for i in 4:-1:1
        _get_col!(m, col, i)
        if sum(norm.(col)) > 0.5
            col_id = i
            break
        end
    end

    for i in 4:-1:1
        _get_row!(m, row, i)
        if sum(norm.(row)) > 0.5
            row_id = i
            break
        end
    end

    col_val = sqrt(sum(col.^2))
    row_val = sqrt(sum(row.^2))

    for i in 0:3

        col_val = col_val * (-1)^(i)
        row_val = row_val * (-1)^(i÷2)

        a, b, c, d = col./col_val
        p, q, r, s = row./row_val

        if col[row_id]/col_val ≈ row_val && row[col_id]/row_val ≈ col_val
            @assert p^2 + q^2 + r^2 + s^2 ≈ 1
            @assert a^2 + b^2 + c^2 + d^2 ≈ 1
            A = [p + q * im s + r * im; -s + r * im p - q * im]
            B = [a + b * im -d + c * im; d + c * im a - b * im]
            return A, B
        end
    end
    error("Algorithm does not work")
end

function decomposeSO4(m::AbstractMatrix{ComplexF64})
    decomposeSO4(real.(m))
end

"""
    compile2qubit(m::AbstractMatrix{ComplexF64}, N, wires=nothing)

compiles an arbitrary U(4) matrix into a quantum circuit. Algorithm taken from
arxiv:quant-ph/0308006, arxiv:quant-ph/0211002
"""
function compile2qubit(m::AbstractMatrix{ComplexF64}, N, wires=nothing)
    cg = CircuitGate[]
    E = 1 / sqrt(2) .* [1 im 0 0; 0 0 im 1; 0 0 im -1; 1 -im 0 0]
    U = E' * m * E

    P2 = U * transpose(U)

    Diag, K_2 = eigen(P2)
    K_2 = real.(K_2)

    # ensure that eigenvectors for degenerate eigenvalues are orthogonal
    if !(det(K_2) ≈ 1) || !(det(K_2) ≈ -1)
        for i in 1:4
            for j in i+1:4
                if norm(dot(K_2[:,j], K_2[:, i])) > 1e-12
                    K_2[:, [i,j]] .= gramm_schmidt!(K_2[:, [i,j]])
                end
            end
        end
    end

    Diag[1] = Diag[1] / det(K_2)^2
    K_2[:,1] = K_2[:,1] * det(K_2)

    Diag = diagm(sqrt.(Diag))
    Diag[1] = Diag[1] * det(U) / det(Diag)

    P = K_2 * Diag * inv(K_2)
    K_1 = inv(P) * U
    K = inv(K_2) * K_1

    C, D = decomposeSO4(inv(K_2) * K_1)

    # println(norm(E'*kron(C,D)*E - inv(K_2) * K_1))
    append!(cg, compile1qubit(C, [2])[1])
    append!(cg, compile1qubit(D, [1])[1])

    append!(cg, [circuit_gate(2, X, 1),
                    circuit_gate((1,), HadamardGate()),
                    circuit_gate((2,), SdagGate()),
                    circuit_gate((1,), SdagGate())])
    append!(cg, compilediag(diag(Diag), N))
    append!(cg, [circuit_gate((1,), SGate()),
                    circuit_gate((2,), SGate()),
                    circuit_gate((1,), HadamardGate()),
                    circuit_gate(2, X, 1)])

    A, B = decomposeSO4(K_2)

    # println(norm(E'*kron(A,B)*E - K_2))
    append!(cg, compile1qubit(A, [2])[1])
    append!(cg, compile1qubit(B, [1])[1])

    return cg
end

"""
    fillψ!(d::Vector{ComplexF64}, ψ::Vector{Float64}, l::Integer)

fills empty vector `ψ`, such that  ψ = −im*[logχ1(d) log χ2(d) ··· logχl−1(U)],
where χn(d) = d[2n-1]*d[2n+2]/(d[2n]*d[2n+1]) per arxiv:quant-ph/0303039
"""
function fillψ!(d::AbstractVector{ComplexF64}, ψ::Vector{Float64}, l::Integer)
    for i in StepRange(1, 1, l - 1)
        ψ[i] = imag(log(d[2i - 1] * d[2i + 2] / (d[2i] * d[2i + 1])))
    end
end

"""
    ηcol(l::Integer, n::Integer)

calculates a single flip state for given integer `n` and a column length of `l`,
    where 0 <= n <= l. returns a vector of length l, where vec[n] = 1, vec[n+1]
    = -1. if n == 0, vec[1] = -1, if n==l, vec[l] = 1
# Examples
```julia-repl
julia> l = 4
julia> ηcol(4, 0)
[-1, 0, 0, 0]
julia> ηcol(4, 1)
[1, -1, 0, 0]
julia> ηcol(4, 2)
[0, 1, -1, 0]
julia> ηcol(4, 3)
[0, 0, 1, -1]
julia> ηcol(4, 4)
[0, 0, 0, 1]
```
"""
@memoize Dict function ηcol(l::Integer, n::Integer)
    vec = zeros(Integer, l)
    if n == 0
        vec[1] = -1
        return vec
    elseif n == l
    vec[l] = 1
        return vec
    end

    vec[n] = 1
    vec[n + 1] = -1
    vec
end

"""
    flip_state(m::Integer, l::Integer)

calculates the `m`th column of the η⊗ matrix by summing all flip states
associated with the binary representation of `m` over the span of [1,l-1]
algorithm taken from arxiv:quant-ph/0303039
"""
@memoize Dict function flip_state(m::Integer, l::Integer)
    # calculates flip states of a given integer `x`
    f = filter(x -> count_ones(x & m) % 2 == 1, 1:l - 1)
    sum(ηcol.((l - 1,), f))
end

"""
    svalue(i::Integer)

returns tuple of position of 1s in the binary representation of integer `i`
starting from least significant digit.
# Examples
```julia-repl
julia> svalue(2)
(2,)
julia> svalue(4)
(3,)
julia> svalue(7)
(1,2,3)
```
"""
function svalue(i::Integer)
    b = Int[]
    count = 1
    while i != 0
        if i & 1 == 1
            push!(b, count)
        end
        count += 1
i = i >> 1
    end
    Tuple(b)
end

"""
    compilediag(d::Vector{ComplexF64}, N, cg=nothing, j=0)

recursively compiles an arbitrary diagonal unitary matrix D_{2^N} ∈ C^(2^N×2^N)
into a quantum circuit. algorithm taken from arxiv:quant-ph/0303039 D_{2^N} can
be decomposed into D_{2^(N-1)} ∈ C^(2^(N-1)×2^(N-1)) ⊗ exp(i*ϕ)*RzGate(θ)
"""
function compilediag(d::AbstractVector{ComplexF64}, N, cg=nothing, j=0)
    cgs = CircuitGate[]
    # single qubit diagonal matrix can be split into a phase shift + RzGate
    if length(d) == 2
        logd = imag.(log.(d))
        β = (logd[1] - logd[2])
        if isnothing(cg)
            cg = CircuitGate[]
        end
        if !(β ≈ 0)
            push!(cg, circuit_gate((j + 1,), RzGate(-β)))
        end
        return cg
    end
    l = 2^(N - j - 1)
    ψ = zeros(Float64, l - 1)
    α = zeros(Float64, l)

    fillψ!(d, ψ, l)

    ηplus = zeros(Float64, (l - 1, l - 1))
    for i in StepRange(1, 1, l - 1)
    ηplus[:,i] = flip_state(greyencode(i), l)
    end

    ηplus = inv(ηplus)
    α[2:end] = -0.5 * ηplus * ψ

    gatewires = svalue(greyencode(1))
    oldwires = ()

    for i in StepRange(2, 1, l)
        rz = [exp(-im * α[i] / 2), exp(im * α[i] / 2)]
        for _ in 1:N - j - 1
            rz = append!(rz, rz)
        end
        s = filter(x -> count_ones(x & greyencode(i - 1)) % 2 == 1, 1:l - 1)
        s = append!(2 .* s .+ 1, 2 .* s .+ 2)

        for k in symdiff(gatewires, oldwires)
            push!(cgs, circuit_gate((j + 1), X, k + 1 + j))
        end
        oldwires = gatewires
        gatewires = svalue(greyencode(i))
        if !(α[i]≈0)
            push!(cgs, circuit_gate((j + 1,), RzGate(α[i])))
        end
        rz[s] .= 1 ./ rz[s]
        d = d ./ rz
    end

    for k in reverse(oldwires)
        push!(cgs, circuit_gate(j + 1, X, k + 1 + j))
    end

    if !(d[1] ≈ d[2])
        logd = imag.(log.(d))
        α[1] = logd[1] - logd[2]
        if !(α[1] ≈ 0)
            pushfirst!(cgs, circuit_gate((j + 1,), RzGate(-α[1])))
        end
        d[1:2:end] = d[1:2:end] ./ exp(im * α[1] / 2)
    end
    if isnothing(cg)
        return compilediag(d[1:2:end], N, cgs, j + 1)
    end
    return compilediag(d[1:2:end], N, append!(cg, cgs), j + 1)
end


@testset "others" begin
    @testset "TuringWishart" begin
        dim = 3
        A = Matrix{Float64}(I, dim, dim)
        dW1 = Wishart(dim + 4, A)
        dW2 = TuringWishart(dim + 4, A)
        dW3 = TuringWishart(dW1)
        mean = Distributions.mean
        @testset "$F" for F in (size, rank, mean, meanlogdet, entropy, cov, var)
            @test F(dW1) == F(dW2) == F(dW3)
        end
        @test Matrix(mode(dW1)) == mode(dW2) == mode(dW3)
        xw = rand(dW2)
        @test insupport(dW1, xw)
        @test insupport(dW2, xw)
        @test insupport(dW3, xw)
        @test logpdf(dW1, xw) == logpdf(dW2, xw) == logpdf(dW3, xw)
    end

    @testset "TuringInverseWishart" begin
        dim = 3
        A = Matrix{Float64}(I, dim, dim)
        dIW1 = InverseWishart(dim + 4, A)
        dIW2 = TuringInverseWishart(dim + 4, A)
        dIW3 = TuringInverseWishart(dIW1)
        mean = Distributions.mean
        @testset "$F" for F in (size, rank, mean, cov, var)
            @test F(dIW1) == F(dIW2) == F(dIW3)
        end
        @test Matrix(mode(dIW1)) == mode(dIW2) == mode(dIW3)
        xiw = rand(dIW2)
        @test insupport(dIW1, xiw)
        @test insupport(dIW2, xiw)
        @test insupport(dIW3, xiw)
        @test logpdf(dIW1, xiw) == logpdf(dIW2, xiw) == logpdf(dIW3, xiw)
    end

    @testset "TuringMvNormal" begin
        @testset "$TD" for TD in [TuringDenseMvNormal, TuringDiagMvNormal, TuringScalMvNormal]
            m = rand(3)
            if TD <: TuringDenseMvNormal
                A = rand(3, 3)
                C = A' * A + I
                d1 = TuringMvNormal(m, C)
                d2 = MvNormal(m, C)
            elseif TD <: TuringDiagMvNormal
                C = rand(3)
                d1 = TuringMvNormal(m, C)
                d2 = MvNormal(m, Diagonal(C .^ 2))
            else
                C = rand()
                d1 = TuringMvNormal(m, C)
                d2 = MvNormal(m, C^2 * I)
            end

            @testset "$F" for F in (length, size, mean)
                @test F(d1) == F(d2)
            end
            @test cov(d1) ≈ cov(d2)
            @test var(d1) ≈ var(d2)

            x1 = rand(d1)
            x2 = rand(d1, 3)
            @test isapprox(logpdf(d1, x1), logpdf(d2, x1), rtol = 1e-6)
            @test isapprox(logpdf(d1, x2), logpdf(d2, x2), rtol = 1e-6)
        end
    end

    @testset "TuringMvLogNormal" begin
        @testset "$TD" for TD in [TuringDenseMvNormal, TuringDiagMvNormal, TuringScalMvNormal]
            m = rand(3)
            if TD <: TuringDenseMvNormal
                C = Matrix{Float64}(I, 3, 3)
                d1 = TuringMvLogNormal(TuringMvNormal(m, C))
                d2 = MvLogNormal(MvNormal(m, C))
            elseif TD <: TuringDiagMvNormal
                C = ones(3)
                d1 = TuringMvLogNormal(TuringMvNormal(m, C))
                d2 = MvLogNormal(MvNormal(m, Diagonal(C .^ 2)))
            else
                C = 1.0
                d1 = TuringMvLogNormal(TuringMvNormal(m, C))
                d2 = MvLogNormal(MvNormal(m, C^2 * I))
            end

            @test length(d1) == length(d2)

            x1 = rand(d1)
            x2 = rand(d1, 3)
            @test logpdf(d1, x1) ≈ logpdf(d2, x1) rtol=1e-6
            @test logpdf(d1, x2) ≈ logpdf(d2, x2) rtol=1e-6

            x2[:, 1] .= -1
            @test isinf(logpdf(d1, x2)[1])
            @test isinf(logpdf(d2, x2)[1])
        end
    end

    @testset "TuringUniform" begin
        @test logpdf(TuringUniform(), 0.5) == 0
    end

    @testset "TuringPoissonBinomial" begin
        d1 = TuringPoissonBinomial([0.5, 0.5])
        d2 = PoissonBinomial([0.5, 0.5])
        @test quantile(d1, 0.5) == quantile(d2, 0.5)
        @test minimum(d1) == minimum(d2)
    end

    @testset "Inverse of pi" begin
        @test 1/pi == inv(pi)
    end

    @testset "Cholesky" begin
        A = rand(3, 3)'; A = A + A' + 3I;
        C = cholesky(A; check = true)
        factors, info = DistributionsAD.turing_chol(A, true)
        @test factors == C.factors
        @test info == C.info
    end

    @testset "zygote_ldiv" begin
        A = to_posdef(rand(3, 3))
        B = copy(A)
        @test DistributionsAD.zygote_ldiv(A, B) == A \ B
    end

    @testset "Entropy" begin
        sigma = exp(randn())
        d1 = TuringScalMvNormal(randn(10), sigma)
        d2 = MvNormal(randn(10), sigma^2 * I)
        @test entropy(d1) ≈ entropy(d2) rtol=1e-6

        sigmas = exp.(randn(10))
        d1 = TuringDiagMvNormal(randn(10), sigmas)
        d2 = MvNormal(randn(10), Diagonal(sigmas .^ 2))
        @test entropy(d1) ≈ entropy(d2) rtol=1e-6

        A = randn(10)
        C = A * A' + I
        d1 = TuringDenseMvNormal(randn(10), C)
        d2 = MvNormal(randn(10), C)
        @test entropy(d1) ≈ entropy(d2) rtol=1e-6
    end

    @testset "Params" begin
        m = rand(10)
        sigmas = randexp(10)

        d = TuringDiagMvNormal(m, sigmas)
        @test params(d) == (m, sigmas)

        d = TuringScalMvNormal(m, sigmas[1])
        @test params(d) == (m, sigmas[1])
    end

    @testset "adapt_randn" begin
        rng = MersenneTwister()
        for T in (Float32, Float64)
            test_adapt_randn(rng, rand(rng, T, 50), T, 10, 30)
        end
    end

    @testset "TuringDirichlet" begin
        dim = 3
        n = 4
        for alpha in (2, rand())
            d1 = TuringDirichlet(dim, alpha)
            d2 = Dirichlet(dim, alpha)
            d3 = TuringDirichlet(d2)
            @test d1.alpha == d2.alpha == d3.alpha
            @test d1.alpha0 == d2.alpha0 == d3.alpha0
            @test d1.lmnB == d2.lmnB == d3.lmnB

            s1 = rand(d1)
            @test s1 isa Vector{Float64}
            @test length(s1) == dim

            s2 = rand(d1, n)
            @test s2 isa Matrix{Float64}
            @test size(s2) == (dim, n)
        end

        for alpha in (ones(Int, dim), rand(dim))
            d1 = TuringDirichlet(alpha)
            d2 = Dirichlet(alpha)
            d3 = TuringDirichlet(d2)
            @test d1.alpha == d2.alpha == d3.alpha
            @test d1.alpha0 == d2.alpha0 == d3.alpha0
            @test d1.lmnB == d2.lmnB == d3.lmnB

            s1 = rand(d1)
            @test s1 isa Vector{Float64}
            @test length(s1) == dim

            s2 = rand(d1, n)
            @test s2 isa Matrix{Float64}
            @test size(s2) == (dim, n)
        end

        # https://github.com/TuringLang/DistributionsAD.jl/issues/158
        let
            d = TuringDirichlet(rand(2))
            z = rand(d)
            logpdf_z = logpdf(d, z)
            pdf_z = pdf(d, z)

            for x in ([0.5, 0.8], [-0.5, 1.5])
                @test logpdf(d, x) == -Inf
                @test iszero(pdf(d, x))

                xmat = hcat(x, x)
                @test all(==(-Inf), logpdf(d, xmat))
                @test all(iszero, pdf(d, xmat))

                xzmat = hcat(x, z)
                @test logpdf(d, xzmat) == [-Inf, logpdf_z]
                @test pdf(d, xzmat) == [0, pdf_z]
            end
        end
    end
end