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

class TeukolskyData(object):

    coeffs = None

    _libtd = None
    _tdctx = None

    class _TDContext(ctypes.Structure):
     _fields_ = [("priv",                   ctypes.c_void_p),
                 ("log_callback",           ctypes.c_void_p),
                 ("opaque",                 ctypes.c_void_p),
                 ("amplitude",              ctypes.c_double),
                 ("nb_coeffs",              ctypes.c_uint * 2),
                 ("basis_scale_factor",     ctypes.c_double * 2),
                 ("max_iter",               ctypes.c_uint),
                 ("atol",                   ctypes.c_double),
                 ("nb_threads",             ctypes.c_uint),
                 ("coeffs",                 ctypes.POINTER(ctypes.c_double) * 3),
                 ("solution_branch",        ctypes.c_uint),
                 ]

    def __init__(self, **kwargs):
        self._libtd = ctypes.CDLL('libteukolskydata.so')
        tdctx_alloc = self._libtd.td_context_alloc
        tdctx_alloc.restype = ctypes.POINTER(self._TDContext)
        self._tdctx = self._libtd.td_context_alloc()

        coeffs_init = ctypes.c_void_p()

        for arg, value in kwargs.iteritems():
            if arg == 'coeffs_init':
                coeffs_init = (ctypes.POINTER(ctypes.c_double) * 3)()
                coeffs_init[0] = ctypes.cast(np.ctypeslib.as_ctypes(value[0]), ctypes.POINTER(ctypes.c_double))
                coeffs_init[1] = ctypes.cast(np.ctypeslib.as_ctypes(value[1]), ctypes.POINTER(ctypes.c_double))
                coeffs_init[2] = ctypes.cast(np.ctypeslib.as_ctypes(value[2]), ctypes.POINTER(ctypes.c_double))
                continue

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

        ret = self._libtd.td_solve(self._tdctx, coeffs_init)
        if ret < 0:
            raise RuntimeError('Error solving the equation')

        self.coeffs = [None] * 3
        for i in xrange(3):
            self.coeffs[i] = np.copy(np.ctypeslib.as_array(self._tdctx.contents.coeffs[i], (self._tdctx.contents.nb_coeffs[1], self._tdctx.contents.nb_coeffs[0])))

    def __del__(self):
        if self._tdctx:
            addr_tdctx = ctypes.c_void_p(ctypes.addressof(self._tdctx))
            self._libtd.td_context_free(addr_tdctx)
            self._tdctx = None

    def _eval_var(self, eval_func, r, theta, diff_order = None):
        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]

        if r.ndim == 2:
            if r.shape != theta.shape:
                raise TypeError('r and theta must be identically-shaped 2-dimensional arrays')
            R, Theta = r.view(), theta.view()
        elif r.ndim == 1:
            if theta.ndim != 1:
                raise TypeError('r and theta must both be 1-dimensional NumPy arrays')
            R, Theta = np.meshgrid(r, theta)
        else:
            raise TypeError('invalid r/theta parameters')

        out = np.empty(R.shape[0] * R.shape[1])

        R.shape     = out.shape
        Theta.shape = out.shape

        c_out   = np.ctypeslib.as_ctypes(out)
        c_r     = np.ctypeslib.as_ctypes(R)
        c_theta = np.ctypeslib.as_ctypes(Theta)

        ret = eval_func(self._tdctx, out.shape[0], c_r, c_theta,
                        c_diff_order, c_out, ctypes.c_long(r.shape[0]))
        if ret < 0:
            raise RuntimeError('Error evaluating the variable')

        out.shape = (theta.shape[0], r.shape[0])
        return out

    def eval_psi(self, r, theta, diff_order = None):
        return self._eval_var(self._libtd.td_eval_psi, r, theta, diff_order)
    def eval_krr(self, r, theta, diff_order = None):
        return self._eval_var(self._libtd.td_eval_krr, r, theta, diff_order)
    def eval_kpp(self, r, theta, diff_order = None):
        return self._eval_var(self._libtd.td_eval_kpp, r, theta, diff_order)


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