# 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/>.
from gmsh import option
import numpy as np


import ctypes
import ctypes.util
import os
from ._tools import timeit
from .csr import CSR

dir_path = os.path.dirname(os.path.realpath(__file__))
lib2 = np.ctypeslib.load_library("libmbfluid2",dir_path)

options_key = {"verbosity":4, "fill_in":14, "reorder_metis":7, "scaling":8, "openmp_threads" :16}

def _np2c(a,dtype=float,order="C") :
    tmp = np.require(a,dtype,order)
    r = ctypes.c_void_p(tmp.ctypes.data)
    r.tmp = tmp
    return r

class LinearSystem2 :
    @timeit
    def __init__(self, elements, n_fields, options=None, constraints=[]) :
        self.localsize = elements.shape[1]*n_fields
        self.constraints = list(c[:,0]*n_fields + c[:,1] for c in constraints)
        idx = ((elements*n_fields)[:,:,None]+np.arange(n_fields)[None,None,:]).reshape([elements.shape[0],-1])
        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)
        row = np.repeat(np.arange(self.csr.size,dtype=np.int32),self.csr.row[1:]-self.csr.row[:-1])
        lib2.mumps_new.restype = ctypes.c_void_p
        self.mumps_ptr = ctypes.c_void_p(lib2.mumps_new(ctypes.c_int(self.csr.size),ctypes.c_int(self.csr.row[-1]),_np2c(row+1,np.int32),_np2c(self.csr.col+1,np.int32)))  
        options = None if options=="" else options
        if options is not None:
            options_vec = []
            for key, val in options.items():
                if isinstance(key,str):
                    key = options_key[key]
                    options_vec.extend([key,val])
            options_vec = np.array(options_vec, dtype=np.int32)
            lib2.mumps_set_options(self.mumps_ptr, _np2c(options_vec, np.int32), ctypes.c_int(np.size(options_vec)//2))

    @timeit
    def local_to_global(self,localv,localm,u, constraints_value):
        self.csr.assemble_mat(localm, constraints_value, self.val)
        return self.csr.assemble_rhs(localv, u, constraints_value)

    @timeit
    def solve(self,rhs) :
        r = np.copy(rhs)
        lib2.mumps_solve(self.mumps_ptr,_np2c(self.val),_np2c(r))
        return r[:self.csr.ndof]

    def __delete__(self):
        lib2.mumps_element_delete(self.mumps_ptr)


class LinearSystem :
    @timeit
    def __init__(self, elements, n_fields, options=None, constraints=[]) :
        n_nodes = np.max(elements)+1
        self.globalsize = n_nodes*n_fields+len(constraints)
        self.localsize = elements.shape[1]*n_fields
        self.idx = (elements[:,None,:]*n_fields+np.arange(n_fields)[None,:,None])
        idxf = [self.idx.swapaxes(1,2)+1]
        eltptr = [np.arange(1,elements.size*n_fields+2,self.localsize)]
        self.ndof = n_nodes*n_fields
        self.n_fields = n_fields
        self.constraints = constraints
        for i,constraint in enumerate(constraints):
            idxc = np.c_[np.full(constraint.shape[0], self.ndof+i+1), constraint[:,0]*n_fields+constraint[:,1]+1].flatten()
            idxf.append(idxc)
            eltptr.append(eltptr[-1][-1]+np.arange(2,constraint.shape[0]*2+2,2))
        lib2.mumps_element_new.restype = ctypes.c_void_p
        eltptr = np.concatenate([i.flat for i in eltptr])
        idxf = np.concatenate([i.flat for i in idxf])
        self.mumps_ptr = ctypes.c_void_p(lib2.mumps_element_new(ctypes.c_int(self.globalsize),ctypes.c_int(eltptr.size-1),_np2c(eltptr,np.int32),_np2c(idxf,np.int32)))  
        options = None if options=="" else options
        if options is not None:
            options_vec = []
            for key, val in options.items():
                if isinstance(key,str):
                    key = options_key[key]
                    options_vec.extend([key,val])
            options_vec = np.array(options_vec, dtype=np.int32)
            lib2.mumps_set_options(self.mumps_ptr, _np2c(options_vec, np.int32), ctypes.c_int(np.size(options_vec)//2))

    @timeit
    def local_to_global(self,localv,localm,u, constraints_values):
        rhs = np.bincount(self.idx.flatten(),localv,self.globalsize)
        for i, (c, cv) in enumerate(zip(self.constraints,constraints_values)):
            rhs[i+self.ndof] = cv[1] + np.sum(u[c[:,0]*self.n_fields+c[:,1]]*cv[0])
        self.v = np.empty(localm.size+sum(c.shape[0]*4 for c in self.constraints))
        p = localm.size
        self.v[:p] = localm.flat
        for constraint, cv  in constraints_values:
            localc = np.zeros((constraint.shape[0],4))
            localc[:,1] = constraint
            localc[:,2] = constraint
            self.v[p:localc.size+p] = localc.flat
            p += localc
        return rhs

    @timeit
    def solve(self,rhs) :
        r = np.copy(rhs)
        lib2.mumps_element_solve(self.mumps_ptr,_np2c(self.v),_np2c(r))
        return r[:self.ndof]

    def __delete__(self):
        lib2.mumps_element_delete(self.mumps_ptr)