import petsc4py
import sys
petsc4py.init(sys.argv)
from petsc4py import PETSc
import numpy as np
from ._tools import timeit
from .csr import CSR

class LinearSystem:

    def __init__(self, elements, n_fields, options, constraints) :
        PETSc.Options().prefixPush("fluid_")
        PETSc.Options().insertString("-pc_type ilu -pc_factor_levels 5 -pc_factor_shift_type INBLOCKS -pc_factor_shift_amount 1e-16")
        PETSc.Options().insertString(options)
        PETSc.Options().prefixPop()
        self.block = PETSc.Options().getBool("fluid_block_matrix",False)
        self.ksp = PETSc.KSP().create()
        self.ksp.setOptionsPrefix(b"fluid_")
        self.ksp.setFromOptions()
        self.localsize = elements.shape[1]*n_fields
        if self.block:
            if len(constraints) != 0 :
                raise ValueError("petsc block matrices not implemented with constraints")
            nn = elements.shape[1]
            self.csr = CSR(elements,elements,[])
            self.mat = PETSc.Mat().createBAIJ(self.csr.size*n_fields, n_fields, csr=(self.csr.row, self.csr.col))
            self.idx = (elements[:,None,:]*n_fields+np.arange(n_fields)[None,:,None]).reshape([-1])
            self.n_fields = n_fields
            csrmapl = np.arange(n_fields*n_fields)[None,:].reshape(n_fields,n_fields)
            csrmap = (self.csr.map[0]*n_fields*n_fields)[:,None,None]+csrmapl[None,:,:]
            self.csrmap = np.swapaxes(csrmap.reshape([elements.shape[0],nn,n_fields,nn,n_fields]),2,3).reshape([-1])
        else:
            idx = ((elements*n_fields)[:,:,None]+np.arange(n_fields)[None,None,:]).reshape([elements.shape[0],-1])
            self.constraints = list(c[:,0]*n_fields + c[:,1] for c in constraints)
            idx2 = (elements[:,None,:]*n_fields+np.arange(n_fields)[None,:,None]).reshape([-1])
            self.csr = CSR(idx, idx2, self.constraints)
            self.val = np.zeros(self.csr.col.size)
            self.mat = PETSc.Mat().createAIJWithArrays(self.csr.size,(self.csr.row,self.csr.col,self.val))

    @timeit
    def local_to_global(self, localv, localm, u, constraints_value):
        if self.block:
            rhs = np.bincount(self.idx,localv,self.csr.size*self.n_fields)
            val = np.bincount(self.csrmap,localm,self.csr.col.size*self.n_fields**2)
            self.mat.zeroEntries()
            self.mat.setValuesBlockedCSR(self.csr.row, self.csr.col, val)
            self.mat.assemble()
            return rhs
        else:
            self.csr.assemble_mat(localm, constraints_value, self.val)
            self.mat.assemble()
            return self.csr.assemble_rhs(localv, u, constraints_value)

    @timeit
    def solve(self,rhs) :
        x = np.ndarray(rhs.shape)
        prhs = PETSc.Vec().createWithArray(rhs.reshape([-1]))
        px = PETSc.Vec().createWithArray(x.reshape([-1]))
        self.ksp.setOperators(self.mat)
        self.ksp.solve(prhs,px)
        if self.block:
            return x
        else:
            return x[:self.csr.ndof]