#
# Copyright 2019 Anton Khirnov <anton@khirnov.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#


import os
import ctypes
import numpy as np

DIFF_COEFF_00 = 0
DIFF_COEFF_01 = 1
DIFF_COEFF_10 = 2
DIFF_COEFF_11 = 3
DIFF_COEFF_02 = 4
DIFF_COEFF_20 = 5

BC_TYPE_FIXVAL  = 0
BC_TYPE_REFLECT = 1
BC_TYPE_FALLOFF = 2

BOUNDARY_0L = 0
BOUNDARY_0U = 1
BOUNDARY_1L = 2
BOUNDARY_1U = 3

ERR_DIVERGE = 0xff00

_doubleptr = ctypes.POINTER(ctypes.c_double)

if ctypes.sizeof(ctypes.c_void_p) == 8:
    _ptrdiff_t = ctypes.c_int64
elif ctypes.sizeof(ctypes.c_void_p) == 4:
    _ptrdiff_t = ctypes.c_int32
else:
    raise NotImplementedError('Unsupported pointer size: %d' %\
                              ctypes.sizeof(ctypes.c_void_p))

def _wrap_array(obj, stride, len):
    arr_np = np.ctypeslib.as_array(obj, (len, stride))
    return arr_np[:, :len]

class SolveError(Exception):
    errno = None
    def __init__(self, errno):
        super().__init__(self, 'Error solving the equation: %s (%d) ' % (os.strerror(errno), errno))
        self.errno = errno

class _MG2DBoundary(ctypes.Structure):
    _fields_ = [
        ("type",        ctypes.c_int),
        ("val",         _doubleptr),
        ("val_stride",  ctypes.c_size_t),
    ]

class _MG2DDCBoundarySpec(ctypes.Structure):
    _fields_ = [
        ("flags",        ctypes.c_int),
        ("val",          _doubleptr),
    ]

class _MG2DDiffCoeffs(ctypes.Structure):
    _fields_ = [
        ("data",        _doubleptr),
        ("stride",      _ptrdiff_t),
        ("boundaries",  _MG2DDCBoundarySpec * 4),
    ]

class _MG2DContext(ctypes.Structure):
    _fields_ = [
        ("priv",                   ctypes.c_void_p),
        ("log_callback",           ctypes.c_void_p),
        ("opaque",                 ctypes.c_void_p),
        ("domain_size",            ctypes.c_size_t),
        ("step",                   ctypes.c_double * 2),
        ("fd_stencil",             ctypes.c_size_t),
        ("boundaries",             ctypes.POINTER(_MG2DBoundary) * 4),
        ("maxiter",                ctypes.c_uint),
        ("tol",                    ctypes.c_double),
        ("nb_cycles",              ctypes.c_uint),
        ("nb_relax_pre",           ctypes.c_uint),
        ("nb_relax_post",          ctypes.c_uint),
        ("u",                      _doubleptr),
        ("u_stride",               _ptrdiff_t),
        ("rhs",                    _doubleptr),
        ("rhs_stride",             _ptrdiff_t),
        #("diff_coeffs",            ctypes.POINTER(_MG2DDiffCoeffs) * 6),
        ("diff_coeffs",            _doubleptr * 6),
        ("diff_coeffs_stride",     _ptrdiff_t),
        ("nb_threads",             ctypes.c_uint),
        ("cfl_factor",             ctypes.c_double),
        ("log_level",              ctypes.c_int),
        ("max_levels",             ctypes.c_uint),
        ("max_exact_size",         ctypes.c_size_t),
    ]

class Boundary:
    type   = None
    lines  = None

    _arr = None

    def __init__(self, val, val_stride, max_stencil, domain_size):
        self._arr = np.ctypeslib.as_array(val, (val_stride * max_stencil,))

        self.lines = []
        for i in range(max_stencil):
            line_offset = i * (val_stride - 1)
            line_len    = domain_size + 2 * i
            self.lines.append(self._arr[line_offset:line_offset + line_len])

class DiffCoeffsBoundarySpec:
    _mg2dctx = None
    _dc_idx  = None
    _bnd_idx = None

    def __init__(self, mg2dctx, dc_idx, bnd_idx):
        self._mg2dctx = mg2dctx
        self._dc_idx  = dc_idx
        self._bnd_idx = bnd_idx

    @property
    def val(self):
        dc = self._mg2dctx.contents.diff_coeffs[self._dc_idx].contents
        return np.ctypeslib.as_array(dc.boundaries[self._bnd_idx].val, (self._mg2dctx.contents.domain_size,))

    @property
    def flags(self):
        dc = self._mg2dctx.contents.diff_coeffs[self._dc_idx].contents
        return dc.boundaries[self._bnd_idx].flags
    @flags.setter
    def flags(self, val):
        dc = self._mg2dctx.contents.diff_coeffs[self._dc_idx].contents
        dc.boundaries[self._bnd_idx].flags = val

class DiffCoeffs:

    boundaries = None

    _mg2dctx = None
    _idx     = None

    def __init__(self, mg2dctx, idx):
        self._mg2dctx = mg2dctx
        self._idx     = idx

        #self.boundaries = []
        #for i in range(4):
        #    self.boundaries.append(DiffCoeffsBoundarySpec(mg2dctx, idx, i))

    @property
    def data(self):
        #dc = self._mg2dctx.contents.diff_coeffs[self._idx].contents
        #return _wrap_array(dc.data, dc.stride, self._mg2dctx.contents.domain_size)
        data = self._mg2dctx.contents.diff_coeffs[self._idx]
        stride = self._mg2dctx.contents.diff_coeffs_stride
        return _wrap_array(data, stride, self._mg2dctx.contents.domain_size)

    def __getitem__(self, key):
        return self.data[key]

    def __setitem__(self, key, val):
        self.data[key] = val

class Solver:
    diff_coeffs = None
    boundaries  = None

    _libmg2d        = None
    _mg2dctx        = None

    def __init__(self, domain_size):
        self._lib_init()

        self._mg2dctx = self._libmg2d.mg2d_solver_alloc(domain_size)
        if self._mg2dctx is None:
            raise MemoryError('Error allocating the solver')

        self.diff_coeffs     = []
        for i in range(len(self._mg2dctx.contents.diff_coeffs)):
            self.diff_coeffs.append(DiffCoeffs(self._mg2dctx, i))

        self.boundaries = []
        for i, bnd in enumerate(self._mg2dctx.contents.boundaries):
            boundary = Boundary(bnd.contents.val, bnd.contents.val_stride,
                                self._libmg2d.mg2d_max_fd_stencil(), domain_size)
            self.boundaries.append(boundary)

    def _lib_init(self):
        self._libmg2d = ctypes.CDLL('libmg2d.so')

        solver_alloc = self._libmg2d.mg2d_solver_alloc
        solver_alloc.argtypes = [ctypes.c_size_t]
        solver_alloc.restype  = ctypes.POINTER(_MG2DContext)

        solver_free = self._libmg2d.mg2d_solver_free
        solver_free.argtypes  = [ctypes.POINTER(ctypes.POINTER(_MG2DContext))]
        solver_free.restype   = None

        solver_solve = self._libmg2d.mg2d_solve
        solver_solve.argtypes = [ctypes.POINTER(_MG2DContext)]
        solver_solve.restype  = ctypes.c_int

        solver_print_stats = self._libmg2d.mg2d_print_stats
        solver_print_stats.argtypes = [ctypes.POINTER(_MG2DContext), ctypes.c_char_p]
        solver_print_stats.restype  = None

        max_fd_stencil = self._libmg2d.mg2d_max_fd_stencil
        max_fd_stencil.restype = ctypes.c_uint

    def __del__(self):
        if self._mg2dctx:
            free = self._libmg2d.mg2d_solver_free
            addr_mg2dctx = ctypes.cast((ctypes.addressof(self._mg2dctx)), free.argtypes[0])
            free(addr_mg2dctx)
            self._mg2dctx = None

    @property
    def u(self):
        return _wrap_array(self._mg2dctx.contents.u, self._mg2dctx.contents.u_stride,
                           self._mg2dctx.contents.domain_size)
    @property
    def rhs(self):
        return _wrap_array(self._mg2dctx.contents.rhs, self._mg2dctx.contents.rhs_stride,
                           self._mg2dctx.contents.domain_size)

    def solve(self, **kwargs):
        for arg, value in kwargs.items():
            try:
                self._mg2dctx.contents.__setattr__(arg, value)
            except TypeError as e:
                # try assigning items of an iterable
                try:
                    for i, it in enumerate(value):
                        self._mg2dctx.contents.__getattribute__(arg)[i] = it
                except:
                    raise e

        for i, bnd in enumerate(self.boundaries):
            if bnd.type is None:
                raise ValueError('Boundary %d not set' % i)

            self._mg2dctx.contents.boundaries[i].contents.type = bnd.type

        ret = self._libmg2d.mg2d_solve(self._mg2dctx)
        if ret < 0:
            raise SolveError(-ret)

    def print_stats(self):
        log_level = self._mg2dctx.contents.log_level
        self._mg2dctx.contents.log_level = 0x40
        self._libmg2d.mg2d_print_stats(self._mg2dctx, ctypes.c_char_p(None))
        self._mg2dctx.contents.log_level = log_level