# MigFlow - Copyright (C) <2010-2018>
# <Universite catholique de Louvain (UCL), Belgium
#  Universite de Montpellier, France>
# 	
# List of the contributors to the development of MigFlow: see AUTHORS file.
# Description and complete License: see LICENSE file.
# 	
# This program (MigFlow) is free software: 
# you can redistribute it and/or modify it under the terms of the GNU Lesser General 
# Public License as published by the Free Software Foundation, either version
# 3 of the License, or (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
# 
# You should have received a copy of the GNU Lesser General Public License
# along with this program (see COPYING and COPYING.LESSER files).  If not, 
# see self.mat<http://www.gnu.org/licenses/>.
import numpy as np
import atexit

from ._tools import timeit, _np2c
from .csr import CSR

import ctypes
import os

petscpath = os.environ["PETSC_DIR"]+"/"+os.environ["PETSC_ARCH"]+"/lib"
os.environ["LD_LIBRARY_PATH"] += ":"+petscpath
lib = np.ctypeslib.load_library("libpetsc",petscpath)

lib.PetscInitialize(None, None, None)
COMM_WORLD = ctypes.c_void_p.in_dll(lib,"PETSC_COMM_WORLD")
#atexit.register(lib.PetscFinalize)


class PetscObj:
    
    def __init__(self, constructor, destructor, *args):
        self._p = ctypes.c_void_p()
        self._destructor = destructor
        lib[constructor](*args,ctypes.byref(self._p))

    def __del__(self):
        lib[self._destructor](ctypes.byref(self._p))

    @property
    def _as_parameter_(self) :
        return self._p


class Vec(PetscObj) :

    def __init__(self, v) :
        super().__init__("VecCreateSeqWithArray","VecDestroy",COMM_WORLD,
                ctypes.c_int(v.shape[-1]), ctypes.c_int(v.size),ctypes.c_void_p(v.ctypes.data))

def insert_options(prefix,options):
    lib.PetscOptionsPrefixPush(None, prefix.encode())
    lib.PetscOptionsInsertString(None, options.encode())
    lib.PetscOptionsPrefixPop(None)

def get_option_bool(prefix, name, default):
    v = ctypes.c_int(0)
    lib.PetscOptionsGetBool(None, prefix.encode(), name.encode(), ctypes.c_int(default), ctypes.byref(v))
    return v.value

class LinearSystem:

    def __init__(self, elements, n_fields, options, constraints = []) :
        insert_options("fluid", "-pc_type ilu -pc_factor_levels 5 -pc_factor_shift_type INBLOCKS -pc_factor_shift_amount 1e-16")
        insert_options("fluid", options)
        self.block = get_option_bool("fluid","-block_matrix",0)
        self.ksp = PetscObj("KSPCreate", "KSPDestroy", COMM_WORLD)
        lib.KSPSetOptionsPrefix(self.ksp,b"fluid")
        lib.KSPSetFromOptions(self.ksp)
        if self.block:
            if len(constraints) != 0 :
                raise ValueError("petsc BAIJ not implemented with constraints")
            nnodes = np.max(elements)+1
            nn = elements.shape[1]
            self.csr = CSR(elements,elements,[])
            self.localsize = elements.shape[1]*n_fields
            self.elements = elements
            self.idx = (self.elements[:,None,:]*n_fields+np.arange(n_fields)[None,:,None]).reshape([-1])
            self.size = nnodes*n_fields
            self.n_fields = n_fields
            csrmapl = np.arange(nn*nn)[None,:].reshape(nn,nn).T
            csrmap = (self.csr.map[0]*nn*nn)[:,None,None]+csrmapl[None,:,:]
            self.csrmap = np.copy(np.swapaxes(csrmap.reshape([elements.shape[0],nn,nn, n_fields,n_fields]),2,3).reshape([-1]))
            self.vsize = np.max(self.csrmap)+1
            self.mat = PetscObj("MatCreate", "MatDestroy", COMM_WORLD)
            lib.MatSetType(self.mat, b"seqbaij")
            ms = self.csr.size*n_fields
            lib.MatSetSizes(self.mat,ctypes.c_int(ms),ctypes.c_int(ms),ctypes.c_int(ms),ctypes.c_int(ms))
            lib.MatSeqBAIJSetPreallocationCSR(self.mat,
                    ctypes.c_int(n_fields), _np2c(self.csr.row, np.int32), _np2c(self.csr.col, np.int32), None)
        else:
            idx = ((elements*n_fields)[:,:,None]+np.arange(n_fields)[None,None,:]).reshape([elements.shape[0],-1])
            self.localsize = elements.shape[1]*n_fields
            self.constraints = list(c[:,0]*n_fields + c[:,1] for c in constraints)
            self.idx = (elements[:,None,:]*n_fields+np.arange(n_fields)[None,:,None]).reshape([-1])
            self.csr = CSR(idx, self.idx,self.constraints)
            self.val = np.zeros(self.csr.col.size)
            self.mat = PetscObj("MatCreateSeqAIJWithArrays", "MatDestroy", COMM_WORLD,
                    ctypes.c_int(self.csr.size), ctypes.c_int(self.csr.size), 
                    _np2c(self.csr.row,np.int32), _np2c(self.csr.col,np.int32), _np2c(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)
            aptr = ctypes.POINTER(ctypes.c_double)()
            lib.MatSeqBAIJGetArray(self.mat, ctypes.byref(aptr))
            a = np.ctypeslib.as_array(aptr,shape=(self.vsize,))
            a[:] = np.bincount(self.csrmap,localm,self.csr.col.size*self.n_fields**2)
            lib.MatSeqBAIJRestoreArray(self.mat, ctypes.byref(aptr))
        else:
            self.csr.assemble_mat(localm, constraints_value, self.val)
            rhs = self.csr.assemble_rhs(localv, u , constraints_value)
        lib.MatAssemblyBegin(self.mat,0)
        lib.MatAssemblyEnd(self.mat,0)
        return rhs

    @timeit
    def solve(self,rhs) :
        x = np.zeros_like(rhs)
        #lib.KSPSetReusePreconditioner(self, 0)
        lib.KSPSetOperators(self.ksp, self.mat, self.mat)
        if self.block:
            lib.KSPSolve(self.ksp, Vec(rhs.reshape(-1,self.n_fields)), Vec(x))
            return x.reshape(rhs.shape)
        else:
            lib.KSPSolve(self.ksp, Vec(rhs.reshape(-1)), Vec(x))
            return x.reshape(rhs.shape)[:self.csr.ndof]