$$\phantom{******** If you see this on the webpage then the browser could not locate *********}$$
$$\phantom{******** jsMath/easy/load.js or the variable root is set wrong in this file *********}$$
$$\newcommand{\vector}[1]{\left[\begin{array}{c} #1 \end{array}\right]}$$ $$\newenvironment{matrix}{\left[\begin{array}{cccccccccc}} {\end{array}\right]}$$ $$\newcommand{\A}{{\cal A}}$$ $$\newcommand{\W}{{\cal W}}$$

#/* Copyright 2008-2013 Research Foundation State University of New York   */

#/* This file is part of QUB Express.                                      */

#/* QUB Express 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.                                    */

#/* QUB Express 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,     */
#/* named LICENSE.txt, in the QUB Express program directory.  If not, see  */
#/* <http://www.gnu.org/licenses/>.                                        */

import numpy
import scipy
import traceback

from qubx.fast.fast_utils import *
import qubx.fast.model

QUBX_IDLZ_METHODS = ["Threshold", "Viterbi", "SKM"]

import ctypes
try:
    if 'windows' in os.platform.system().lower():
        ctypes.cdll.LoadLibrary('OpenCL.dll')
    maxill = ctypes.cdll.LoadLibrary('maxill_opencl.dll')
except:
    try:
        maxill = ctypes.cdll.LoadLibrary('maxill.dll')
    except:
        try:
            maxill = ctypes.cdll.LoadLibrary('libmaxill_opencl.so')
        except:
            try:
                maxill = ctypes.cdll.LoadLibrary('@executable_path/../Frameworks/libmaxill_opencl.so')
            except:
                try:
                    maxill = ctypes.cdll.LoadLibrary('libmaxill.so')
                except OSError:
                    maxill = ctypes.cdll.LoadLibrary('@executable_path/../Frameworks/libmaxill.so')
    
maxill.qubx_idlz_data_create.argtypes = (c_double, c_int, c_int)
maxill.qubx_idlz_data_create.restype = c_void_p
maxill.qubx_idlz_data_reset.argtypes = (c_void_p, c_double, c_int, c_int)
maxill.qubx_idlz_data_reset.restype = None
maxill.qubx_idlz_data_free.argtypes = (c_void_p,)
maxill.qubx_idlz_data_free.restype = None
maxill.qubx_idlz_data_append_segment.argtypes = (c_void_p, c_int)
maxill.qubx_idlz_data_append_segment.restype = c_int
maxill.qubx_idlz_data_append_chunk.argtypes = (c_void_p, c_int, c_int, c_float_p,
                                              c_int_p, c_int_pp, c_float_pp, c_double_pp)
maxill.qubx_idlz_data_append_chunk.restype = None
maxill.qubx_idlz_data_get_stimclasses.argtypes = (c_void_p, c_int_p)
maxill.qubx_idlz_data_get_stimclasses.restype = c_double_p
maxill.qubx_idlz_data_get_stimclass_frac.argtypes = (c_void_p,)
maxill.qubx_idlz_data_get_stimclass_frac.restype = c_double_p
maxill.qubx_idlz_data_get_dwellcount.argtypes = (c_void_p, c_int)
maxill.qubx_idlz_data_get_dwellcount.restype = c_int
maxill.qubx_idlz_data_get_dwells.argtypes = (c_void_p, c_int, c_int_p, c_int_p, c_int_p, c_double_p, c_double_p)
maxill.qubx_idlz_data_get_dwells.restype = c_void_p
maxill.qubx_idlz_data_get_baseline.argtypes = (c_void_p, c_int, c_int, c_int, c_float_p)
maxill.qubx_idlz_data_get_baseline.restype = c_int

class Data(object):
    def __init__(self, sampling, Nstimsig, data_cap=2**16):
        self.obj = maxill.qubx_idlz_data_create(sampling, Nstimsig, data_cap)
        self.sampling = sampling
        self.Nstimsig = Nstimsig
        self.Nsig = Nstimsig + 1
        self.Ndata = 0
        self.seg_Ndata = []
    def __del__(self):
        maxill.qubx_idlz_data_free(self.obj)
    def reset(self, sampling, Nstimsig, data_cap=2**16):
        maxill.qubx_idlz_data_reset(self.obj, sampling, Nstimsig, data_cap)
        self.sampling = sampling
        self.Nstimsig = Nstimsig
        self.Ndata = 0
    def append_segment(self, first):
        maxill.qubx_idlz_data_append_segment(self.obj, first)
        self.seg_Ndata.append(0)
    def append_chunk(self, seg_ix, Nd, I, dwellCounts, classeses, durationses, ampses):
        dcs = numpy.array(dwellCounts, dtype='int32')
        # assuming the elements of classeses and durationses are numpy.arrays of int, float, double respectively
        ccs = (c_int_p*self.Nstimsig)(*[cdata(classes, c_int_p) for classes in classeses])
        dds = (c_float_p*self.Nstimsig)(*[cdata(durations, c_float_p) for durations in durationses])
        aas = (c_double_p*self.Nstimsig)(*[cdata(amps, c_double_p) for amps in ampses])
        maxill.qubx_idlz_data_append_chunk(self.obj, seg_ix, Nd, cdata(I, c_float_p), 
                                           cdata(dcs, c_int_p), ccs, dds, aas)
        if I != None:
            self.Ndata += Nd
            self.seg_Ndata[-1] += Nd
    def get_stimclasses(self):
        Nstim = c_int()
        self.stimclasses = maxill.qubx_idlz_data_get_stimclasses(self.obj, byref(Nstim))
        self.Nstim = Nstim.value
        #self.stimclasses = numpy.frombuffer(pybuf(stimcls_p, self.Nstim*self.Nsig*sizeof(c_double)), dtype='float64', count=self.Nstim*self.Nsig).reshape((self.Nstim,self.Nsig))
        return self.stimclasses
    def get_stimclass_frac(self):
        self.sc_frac = maxill.qubx_idlz_data_get_stimclass_frac(self.obj)
        #self.sc_frac = numpy.frombuffer(pybuf(sc_frac_p, self.Nstim*sizeof(c_double)), dtype='float64', count=self.Nstim)
        return self.sc_frac
    def get_dwells(self, iseg):
        ncls = c_int()
        count = maxill.qubx_idlz_data_get_dwellcount(self.obj, iseg, byref(ncls))
        if count:
            ff = numpy.zeros(dtype='int32', shape=(count,))
            ll = numpy.zeros(dtype='int32', shape=(count,))
            cc = numpy.zeros(dtype='int32', shape=(count,))
            aa = numpy.zeros(dtype='float64', shape=(ncls.value,))
            ss = numpy.zeros(dtype='float64', shape=(ncls.value,))
            maxill.qubx_idlz_data_get_dwells(self.obj, iseg, cdata(ff, c_int_p), cdata(ll, c_int_p), cdata(cc, c_int_p),
                                             cdata(aa, c_double_p), cdata(ss, c_double_p))
        else:
            ff = ll = cc = numpy.zeros(dtype='int32', shape=(0,))
            aa = ss = numpy.zeros(dtype='float64', shape=(0,))
        return ff, ll, cc, aa, ss
    def get_baseline(self, iseg, offset=0, count=None):
        n = self.seg_Ndata[iseg] if (count is None) else count
        baseline = numpy.zeros(dtype='float32', shape=(n,))
        maxill.qubx_idlz_data_get_baseline(self.obj, iseg, offset, n, cdata(baseline, c_float_p))
        return baseline


maxill.qubx_idlz_halfamp.argtypes = (qubx.fast.model.PModel, c_void_p, qubx.fast.model.PStimAmps)
maxill.qubx_idlz_halfamp.restype = c_int

def halfamp(model, data, ampm):
    model_obj = model.obj if isinstance(model, qubx.fast.model.Model) else model
    result = maxill.qubx_idlz_halfamp(model_obj, data.obj, ampm.obj)
    if (0 != result):
        raise Exception('qubx_idealize.halfamp error code: %d' % result)
    return 0.0

maxill.qubx_idlz_skm.argtypes = (qubx.fast.model.PModel, c_void_p, qubx.fast.model.PStimAmps, qubx.fast.model.PStimRates, c_double, c_int, c_int, c_int, c_int, c_double, c_double_p)
maxill.qubx_idlz_skm.restype = c_int

def skm(model, data, ampm, ratesm, resolution=5e-2, method=1, reest_amps=True, reest_rates=False, track_baseline=False, baseline_std=1e-3):
    model_obj = model.obj if isinstance(model, qubx.fast.model.Model) else model
    ll = c_double()
    result = maxill.qubx_idlz_skm(model_obj, data.obj, ampm.obj, ratesm.obj, resolution, int(method),
                                  int(reest_amps), int(reest_rates), int(track_baseline), baseline_std, byref(ll))
    if (0 != result) and (-10 != result):
        raise Exception('qubx_idealize.skm error code: %d' % result)
    return ll.value