using Pkg 
Pkg.activate("PartitionedSolvers")
using PartitionedArrays
using PartitionedSolvers
using LinearAlgebra
using DelimitedFiles
using IterativeSolvers
using IterativeSolvers: cg! 
using Statistics
using SparseArrays
using Plots
using Printf

# Initialize arguments for A and B
parts_per_dir = (2,2,1)
block_size = length(parts_per_dir) 
p = prod(parts_per_dir)
ranks = DebugArray(LinearIndices((p,)))

# Initialize solvers
level_params = amg_level_params_linear_elasticity(block_size;
    coarsening = PartitionedSolvers.smoothed_aggregation_with_block_size(;)
)
fine_params = amg_fine_params(;level_params)
solver_le = amg(;fine_params)
solver_sc = amg()
solvers = Dict("le" => solver_le, "sc" => solver_sc)

# Initialize data structures for outcomes
problem_sizes = [3,6,9,12,15,18] 
matrix_sizes = zeros(Int, length(problem_sizes))
iters = zeros(Int, (length(solvers), length(matrix_sizes)))
runtime = zeros((length(solvers), length(problem_sizes)))
resnorm = similar(runtime)

for (k,nnodes) in enumerate(problem_sizes)
    # Construct matrices
    nodes_per_dir = map(i->nnodes*i,parts_per_dir)
    args = PartitionedArrays.linear_elasticity_fem(nodes_per_dir,parts_per_dir,ranks)
    A_dist = psparse(args...) |> fetch 
    A_seq = centralize(A_dist)
    matrix_sizes[k] = A_seq.m
    println("Size: $(A_seq.m)")
    x = node_coordinates_unit_cube(nodes_per_dir,parts_per_dir, ranks)
    B_dist = nullspace_linear_elasticity(x, partition(axes(A_dist,2)))
    B_seq = [collect(b) for b in B_dist]
    x_exact = ones(A_seq.n)
    b = A_seq*x_exact   
    y = similar(x_exact)
    
    for (j,key) in enumerate(keys(solvers))
        y .= 0
        @show key
        start = time()
        if key == "le"
            Pl = setup(solvers[key],y,A_seq,b,nullspace=B_seq)
        else 
            Pl = setup(solvers[key],y,A_seq,b)
        end
        elapsed = time() - start
        @show elapsed
        _, history = cg!(y,A_seq,b;Pl,log=true)
        iters[j,k] = history.iters
        runtime[j,k] = elapsed + eps(eltype(elapsed))
        resnorm[j,k] = history[:resnorm][end]
    end
end

yticks = 2:2:maximum(iters)
yticklabels = [ @sprintf("%5.f",y) for y in yticks ]
xticklabels = [ @sprintf("%7.f",y) for y in matrix_sizes ]
println("matrix sizes: $matrix_sizes, iters': $(iters')")
p1 = plot(matrix_sizes, iters', label=["vector valued (D=$block_size)" "scalar"], 
    markershape = :circle, xaxis=:log10,
    xlims=(0,maximum(matrix_sizes)+maximum(matrix_sizes)*0.05), ylims=(0,maximum(iters)+1), 
    xticks=(matrix_sizes, xticklabels), 
    yticks=(yticks,yticklabels),
    ylabel = "# iterations", xlabel="# rows of A (log)",
    xrotation = 60,
    size=(600,470),
    legend=:bottomright,
    title="Convergence of cg (sequential)")
display(p1)
savefig(p1, "C:/Users/gelie/Home/ComputationalScience/GSoC/convergence_seq_amg_$(maximum(matrix_sizes))_error.png")   

p2 = plot(matrix_sizes, runtime', label=["vector valued (D=$block_size)" "scalar"], 
    markershape = :circle, xaxis=:log10, yaxis=:log10,
    xlims=(0,maximum(matrix_sizes)+maximum(matrix_sizes)*0.05), #ylims=(0,maximum(iters)+1), 
    xticks=(matrix_sizes, xticklabels), 
    #yticks=(yticks,yticklabels),
    ylabel = "log runtime (s)", xlabel="# rows of A (log)",
    xrotation = 60,
    size=(600,470),
    legend=:bottomright,
    title="Runtime of setup() (sequential)")
display(p2)
savefig(p2, "C:/Users/gelie/Home/ComputationalScience/GSoC/runtime_seq_setup_$(maximum(matrix_sizes)).png")   

p3 = plot(matrix_sizes, resnorm', label=["vector valued (D=$block_size)" "scalar"], 
    markershape = :circle, xaxis=:log10, yaxis=:log10,
    xlims=(0,maximum(matrix_sizes)+maximum(matrix_sizes)*0.05), #ylims=(0,maximum(iters)+1), 
    xticks=(matrix_sizes, xticklabels), 
    #yticks=(yticks,yticklabels),
    ylabel = "residual norm (log)", xlabel="# rows of A (log)",
    xrotation = 60,
    size=(600,470),
    legend=:bottomright,
    title="Residual norm in final iteration (sequential)")
display(p3)
savefig(p3, "C:/Users/gelie/Home/ComputationalScience/GSoC/resnorm_seq_$(maximum(matrix_sizes)).png")