#
# Copyright 2014 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 brill_data
import ctypes
import numpy as np

BD_METRIC_COMPONENT_RHORHO	= 0
BD_METRIC_COMPONENT_RHOZ	= 1
BD_METRIC_COMPONENT_RHOPHI	= 2
BD_METRIC_COMPONENT_ZRHO	= 3
BD_METRIC_COMPONENT_ZZ	        = 4
BD_METRIC_COMPONENT_ZPHI	= 5
BD_METRIC_COMPONENT_PHIRHO	= 6
BD_METRIC_COMPONENT_PHIZ	= 7
BD_METRIC_COMPONENT_PHIPHI	= 8

class BrillData(object):

    coeffs = None

    _libbd = None
    _bdctx = None

    class _BDContext(ctypes.Structure):
     _fields_ = [("priv",                   ctypes.c_void_p),
                 ("log_callback",           ctypes.c_void_p),
                 ("opaque",                 ctypes.c_void_p),
                 ("q_func_type",            ctypes.c_int),
                 ("amplitude",              ctypes.c_double),
                 ("eppley_n",               ctypes.c_uint),
                 ("rho0",                   ctypes.c_double),
                 ("nb_coeffs",              ctypes.c_uint * 2),
                 ("basis_scale_factor",     ctypes.c_double * 2),
                 ("psi_minus1_coeffs",      ctypes.POINTER(ctypes.c_double)),
                 ("stride",                 ctypes.c_uint),
                 ("nb_threads",             ctypes.c_uint)]

    def __init__(self, **kwargs):
        self._libbd = ctypes.CDLL('libbrilldata.so')
        bdctx_alloc = self._libbd.bd_context_alloc
        bdctx_alloc.restype = ctypes.POINTER(self._BDContext)
        self._bdctx = self._libbd.bd_context_alloc()
        for arg, value in kwargs.items():
            try:
                self._bdctx.contents.__setattr__(arg, value)
            except TypeError as e:
                # try assigning items of an iterable
                try:
                    for i, it in enumerate(value):
                        self._bdctx.contents.__getattribute__(arg)[i] = it
                except:
                    raise e

        ret = self._libbd.bd_solve(self._bdctx)
        if ret < 0:
            raise RuntimeError('Error solving the equation')

        self.coeffs = np.copy(np.ctypeslib.as_array(self._bdctx.contents.psi_minus1_coeffs, (self._bdctx.contents.nb_coeffs[1], self._bdctx.contents.nb_coeffs[0])))

    def __del__(self):
        if self._bdctx:
            addr_bdctx = ctypes.c_void_p(ctypes.addressof(self._bdctx))
            self._libbd.bd_context_free(addr_bdctx)
            self._bdctx = None

    def eval_psi(self, rho, z, diff_order = None):
        if rho.ndim != 1 or z.ndim != 1:
            raise TypeError('rho and z must be 1-dimensional NumPy arrays')

        if diff_order is None:
            diff_order = [0, 0]

        c_diff_order = (ctypes.c_uint * 2)()
        c_diff_order[0] = diff_order[0]
        c_diff_order[1] = diff_order[1]

        psi = np.empty((z.shape[0], rho.shape[0]))

        c_psi = np.ctypeslib.as_ctypes(psi)
        c_rho = np.ctypeslib.as_ctypes(rho)
        c_z   = np.ctypeslib.as_ctypes(z)

        ret = self._libbd.bd_eval_psi(self._bdctx, c_rho, len(c_rho),
                                      c_z, len(c_z), c_diff_order, c_psi, len(c_rho))
        if ret < 0:
            raise RuntimeError('Error evaluating psi')

        return psi

    def eval_metric(self, rho, z, component, diff_order = None):
        if rho.ndim != 1 or z.ndim != 1:
            raise TypeError('rho and z must be 1-dimensional NumPy arrays')

        if diff_order is None:
            diff_order = [0, 0]

        c_diff_order = (ctypes.c_uint * 2)()
        c_diff_order[0] = diff_order[0]
        c_diff_order[1] = diff_order[1]

        c_component = (ctypes.c_uint * 1)()
        c_component[0] = component

        c_rho    = np.ctypeslib.as_ctypes(rho)
        c_z      = np.ctypeslib.as_ctypes(z)

        metric = np.empty((z.shape[0], rho.shape[0]))

        c_metric    = (ctypes.POINTER(ctypes.c_double) * 1)()
        c_metric[0] = ctypes.cast(np.ctypeslib.as_ctypes(metric), type(c_metric[0]))

        c_metric_strides = (ctypes.c_uint * 1)()
        c_metric_strides[0] = len(c_rho)

        ret = self._libbd.bd_eval_metric(self._bdctx, c_rho, len(c_rho),
                                         c_z, len(c_z), 1, c_component, c_diff_order, c_metric, c_metric_strides)
        if ret < 0:
            raise RuntimeError('Error evaluating the metric')

        return metric

    def eval_q(self, rho, z, diff_order = None):
        if rho.ndim != 1 or z.ndim != 1:
            raise TypeError('rho and z must be 1-dimensional NumPy arrays')

        if diff_order is None:
            diff_order = [0, 0]

        c_diff_order = (ctypes.c_uint * 2)()
        c_diff_order[0] = diff_order[0]
        c_diff_order[1] = diff_order[1]

        q = np.empty((z.shape[0], rho.shape[0]))

        c_q   = np.ctypeslib.as_ctypes(q)
        c_rho = np.ctypeslib.as_ctypes(rho)
        c_z   = np.ctypeslib.as_ctypes(z)

        ret = self._libbd.bd_eval_q(self._bdctx, c_rho, len(c_rho),
                                    c_z, len(c_z), c_diff_order, c_q, len(c_rho))
        if ret < 0:
            raise RuntimeError('Error evaluating q')

        return q

    @property
    def amplitude(self):
        return self._bdctx.contents.__getattribute__('amplitude')