import pyamgx
import pyamg
from scipy import sparse
import numpy as np
import time
from scipy.sparse import spdiags,linalg,eye

pyamgx.initialize()

convergence_limit_p=1e-10
max_amg_iters = 2 #max iters for the amg preconditioners
amg_cycle='V'
amg_relax='fgmres'

cfg= pyamgx.Config().create_from_dict({
       "config_version": 2,
       "solver": {
           "print_grid_stats": 0,
           "store_res_history": 1,
           "solver": "FGMRES",
           "print_solve_stats": 0,
           "obtain_timings": 0,
           "preconditioner": {
               "interpolator": "D2",
               "solver": "AMG",
               "smoother": {
                   "scope": "jacobi",
                   "solver": "JACOBI_L1"
               },
               "presweeps": 1,
               "max_iters":max_amg_iters,
               "max_levels": 24,
               "cycle": amg_cycle,
               "postsweeps": 1
           },
           "max_iters": 100,
           "monitor_residual": 1,
           "gmres_n_restart": 10,
           "convergence": "RELATIVE_INI_CORE",
           "tolerance": convergence_limit_p,
           "norm": "L2"
       }
})


rsc = pyamgx.Resources().create_simple(cfg)

# Create matrices and vectors:
A_x = pyamgx.Matrix().create(rsc)
b_x = pyamgx.Vector().create(rsc)
x_x = pyamgx.Vector().create(rsc)

# Create solver:
solver = pyamgx.Solver().create(rsc, cfg)

# load matrix
#su=np.load('/chalmers/users/lada/pythons-rans-code/boundary-layer-IDDES-fk-0-ni-500-nk64-a_synt-from-uin-L-0.2-7500-15000-commut-with-fk-psi-iter-2-to-create-A-and-Su-from-p/su_p_saved.npy')
#phi=np.load('/chalmers/users/lada/pythons-rans-code/boundary-layer-IDDES-fk-0-ni-500-nk64-a_synt-from-uin-L-0.2-7500-15000-commut-with-fk-psi-iter-2-to-create-A-and-Su-from-p/p_saved.npy')
#A = sparse.load_npz('/chalmers/users/lada/pythons-rans-code/boundary-layer-IDDES-fk-0-ni-500-nk64-a_synt-from-uin-L-0.2-7500-15000-commut-with-fk-psi-iter-2-to-create-A-and-Su-from-p/Ap_saved.npz')


su=np.load('su_p_saved.npy')
phi=np.load('p_saved.npy')
A = sparse.load_npz('Ap_saved.npz')

phi_init=phi


A_x.upload_CSR(A) #upload pyamg poisson problem
b_x.upload(su)     #upload pyamg rand rhs
x_x.upload(phi)

# Setup and solve system:
solver.setup(A_x)

t0 = time.time()

solver.solve(b_x, x_x)

t_pyamgx = time.time()-t0

print('pyamgx time: ',t_pyamgx)

# Download solution
x_x.download(phi)

print('pyamgx, p[0:20]: ',phi[0:20])


t1 = time.time()
A = pyamg.ruge_stuben_solver(A)

phi = A.solve(su, tol=convergence_limit_p, x0=phi_init,accel=amg_relax,cycle=amg_cycle)

t_pyamg = time.time()-t1

print('pyamg time: ',t_pyamg)

print('pyamg, p[0:20]: ',phi[0:20])