$$\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.model import *

try:
    importer = ctypes.windll
    base, ext = '','.dll'
except:
    importer = ctypes.cdll
    base, ext = 'lib','.so'

try:
    if ext == '.dll':
        importer.LoadLibrary(base+'OpenCL'+ext) # windows complains if dll's sub-dependencies are missing
    maxmll = importer.LoadLibrary(base+'maxmll_opencl'+ext)
except:
    try:
        maxmll = importer.LoadLibrary(base+'maxmll'+ext)
    except OSError:
        maxmll = importer.LoadLibrary('@executable_path/../Frameworks/libmaxmll_opencl.so') # , mode=ctypes.RTLD_GLOBAL) to be visible from subsequently-loaded libs

class SampleStruct(Structure):
    _fields_ = [('time', c_float),
                ('cur', c_float),
                ('var', c_float)]
PSample = POINTER(SampleStruct)    

class SampleOutStruct(Structure):
    _fields_ = [('fit', c_float),
                ('fit_var', c_float),
                ('ll', c_int)]
PSampleOut = POINTER(SampleOutStruct)    

maxmll.qubx_mac_data_create.argtypes = (c_int, c_int)
maxmll.qubx_mac_data_create.restype = c_void_p
maxmll.qubx_mac_data_reset.argtypes = (c_void_p, c_int, c_int)
maxmll.qubx_mac_data_reset.restype = None
maxmll.qubx_mac_data_free.argtypes = (c_void_p,)
maxmll.qubx_mac_data_free.restype = None
maxmll.qubx_mac_data_append_segment.argtypes = (c_void_p,)
maxmll.qubx_mac_data_append_segment.restype = c_int
maxmll.qubx_mac_data_append_chunk.argtypes = (c_void_p, c_int, c_double, c_int, c_float_p, c_float_p, c_float_p,
                                              c_int_p, c_int_pp, c_float_pp, c_double_pp)
maxmll.qubx_mac_data_append_chunk.restype = None
maxmll.qubx_mac_data_get_stimclasses.argtypes = (c_void_p, c_int_p)
maxmll.qubx_mac_data_get_stimclasses.restype = c_double_p
maxmll.qubx_mac_data_get_stimclass_frac.argtypes = (c_void_p,)
maxmll.qubx_mac_data_get_stimclass_frac.restype = c_double_p
maxmll.qubx_mac_data_get_buffer.argtypes = (c_void_p,)
maxmll.qubx_mac_data_get_buffer.restype = PSampleOut

class Data(object):
    def __init__(self, Nstimsig, data_cap=2**16, file_index=0):
        self.obj = maxmll.qubx_mac_data_create(Nstimsig, data_cap)
        self.Nstimsig = Nstimsig
        self.file_index = file_index
        self.Nsig = Nstimsig + 1
        self.Ndata = 0
    def __del__(self):
        maxmll.qubx_mac_data_free(self.obj)
    def reset(self, Nstimsig, data_cap=2**16):
        maxmll.qubx_mac_data_reset(self.obj, Nstimsig, data_cap)
        self.Nstimsig = Nstimsig
        self.Ndata = 0
    def append_segment(self):
        maxmll.qubx_mac_data_append_segment(self.obj)
    def append_chunk(self, seg_ix, dur_sec, Nd, T, I, V, 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])
        maxmll.qubx_mac_data_append_chunk(self.obj, seg_ix, dur_sec, Nd, cdata(T, c_float_p), cdata(I, c_float_p), 
                                          cdata(V, c_float_p), cdata(dcs, c_int_p), ccs, dds, aas)
        self.Ndata += Nd
    def get_stimclasses(self):
        Nstim = c_int()
        self.stimclasses = maxmll.qubx_mac_data_get_stimclasses(self.obj, byref(Nstim))
        self.Nstim = Nstim.value
        return self.stimclasses
    def get_stimclass_frac(self):
        self.sc_frac = maxmll.qubx_mac_data_get_stimclass_frac(self.obj)
        return self.sc_frac
    def get_buffer(self):
        return numpy.frombuffer(pybuf(maxmll.qubx_mac_data_get_buffer(self.obj), int(3*self.Ndata*sizeof(c_float))),
                                dtype='float32', count=3*self.Ndata).reshape((self.Ndata, 3))


class ParamStruct(Structure):
    _fields_ = [('storage', c_void_p),
                ('data', c_void_p),
                ('useVar', c_int),
                ('optNchannel', c_int),
                ('accel', c_int),
                ('Npar', c_int),
                ('NfPar', c_int),
                ('pars', c_double_p),
                ('fPars', c_double_p),
                ('gpu_fail', c_int),
                ('ll', c_double),
                ('float_Nchannel', c_double)]
PParam = POINTER(ParamStruct)

maxmll.qubx_mac_param_create.argtypes = (c_int,)
maxmll.qubx_mac_param_create.restype = PParam
maxmll.qubx_mac_param_free.argtypes = (PParam,)
maxmll.qubx_mac_param_free.restype = None
maxmll.qubx_mac_param_add_model.argtypes = (PParam, PModel, PStimAmps, PStimRates)
maxmll.qubx_mac_param_add_model.restype = None
maxmll.qubx_mac_setup_constraints.argtypes = (PParam,)
maxmll.qubx_mac_setup_constraints.restype = c_int
maxmll.qubx_mac_do_fPars_to_model.argtypes = (PParam,)
maxmll.qubx_mac_do_fPars_to_model.restype = None
maxmll.qubx_mac_calc_std.argtypes = (PParam, c_double_p)
maxmll.qubx_mac_calc_std.restype = None

class Param(object):
    def __init__(self, Nmodel):
        self.obj = maxmll.qubx_mac_param_create(Nmodel)
        self.pars_ = self.fPars_ = []
        self.__data = None
        self.models = []
        self.ampms = []
        self.ratesms = []
    def set_data(self, x):
        self.obj[0].data = x.obj
        self.__data = x
    def add_model(self, fastmodel, ampm, ratesm):
        maxmll.qubx_mac_param_add_model(self.obj, fastmodel.obj, ampm.obj, ratesm.obj)
        self.models.append(fastmodel)
        self.ampms.append(ampm)
        self.ratesms.append(ratesm)
    data = property(lambda self: self.__data, lambda self, x: self.set_data(x))
    def set_useVar(self, x):
        self.obj[0].useVar = x
    useVar = property(lambda self: self.obj[0].useVar, lambda self, x: self.set_useVar(x))
    def set_optNchannel(self, x):
        self.obj[0].optNchannel = x
    optNchannel = property(lambda self: self.obj[0].optNchannel, lambda self, x: self.set_optNchannel(x))
    def set_accel(self, x):
        self.obj[0].accel = x
    accel = property(lambda self: self.obj[0].accel, lambda self, x: self.set_accel(x))
    def set_float_Nchannel(self, x):
        self.obj[0].float_Nchannel = x
    float_Nchannel = property(lambda self: self.obj[0].float_Nchannel, lambda self, x: self.set_float_Nchannel(x))
    Npar = property(lambda self: self.obj[0].Npar)
    NfPar = property(lambda self: self.obj[0].NfPar)
    pars = property(lambda self: self.pars_)
    fPars = property(lambda self: self.fPars_)
    ll = property(lambda self: self.obj[0].ll)
    gpu_fail = property(lambda self: self.obj[0].gpu_fail)
    def __del__(self):
        maxmll.qubx_mac_param_free(self.obj)
    def setup_constraints(self):
        result = maxmll.qubx_mac_setup_constraints(self.obj)
        self.pars_ = numpy.frombuffer(pybuf(self.obj[0].pars, self.Npar*sizeof(c_double)), dtype='float64', count=self.Npar)
        self.fPars_ = numpy.frombuffer(pybuf(self.obj[0].fPars, max(1, self.NfPar)*sizeof(c_double)), dtype='float64', count=max(1, self.NfPar))
        return result
    def do_fPars_to_model(self):
        maxmll.qubx_mac_do_fPars_to_model(self.obj)
    def calc_std(self, InvHessian):
        hess = numpy.array(InvHessian, dtype='float64', copy=True)
        maxmll.qubx_mac_calc_std(self.obj, cdata(hess, c_double_p))
        


maxmll.qubx_mac_ll.argtypes = (PParam,)
maxmll.qubx_mac_ll.restype = c_int

def mac_ll(param):
    if param.data.Ndata == 0:
        return 0.0
    ll = c_double()
    result = maxmll.qubx_mac_ll(param.obj)
    if (0 != result) and (-8 != result):
        raise Exception('Mac error code: %d' % result)
    return param.ll



def setup_test():
    def pr(x):
        print x
    array = numpy.array

    data = Data(1) # TODO: test 0 stim signals
    data.append_segment()
    data.append_chunk(0, .0023, 0, None, None, None, [1], [array([1])], [array([2.3], dtype='float32')], [array([-40.0, 10.0])])
    data.append_chunk(0, .001, 1, array([.003], dtype='float32'), array([50.0], dtype='float32'), array([0.0], dtype='float32'),
                      [1], [array([0])], [array([1.0], dtype='float32')], [array([-40.0, 10.0])])
    model = Model(2, 2, [0, 1], pr)
    model.usePeq = 1
    model.useVar = 1
    model.IeqFv = 1
    model.amp0 = 0.0
    model.var0 = 0.01
    model.Nchannel = 100
    model.iVoltage = 1
    model.Gleak = 0.0
    model.Vrev = -80.0
    model.amp[1] = 10.0
    model.var[1] = 1.0
    model.K0[0,1] = 1000.0
    model.K0[1,0] = 1000.0
    model.K1[0,1] = 0.1
    model.V[0,1] = 1
    model.setup_stim(data)
    model.fix_channelcount()
    model.setup_constraints()
    
    return model, data

def setup_data(data_path, delta):
    file = qubx.data_types.Open(data_path)
    
    for f,l,n in file.segmentation.segments:
        pass ###
    data = Data(Nstimsig) # TODO: test 0 stim signals
    data.append_segment()
    data.append_chunk(0, .0023, 0, None, None, None, [1], [array([1])], [array([2.3], dtype='float32')], [array([-40.0, 10.0])])
    data.append_chunk(0, .001, 1, array([.003], dtype='float32'), array([50.0], dtype='float32'), array([0.0], dtype='float32'),
                      [1], [array([0])], [array([1.0], dtype='float32')], [array([-40.0, 10.0])])

def setup_model(model_path, data):
    def pr(x):
        print x
    array = numpy.array
    
    model, model_file = setup_model(model_path)
    file = qubx.model.QubModel()
    file.from_tree(qubx.tree.Open(model_path, True))
    clazz = ModelToClazz(file)
    Nclass = int(max(clazz) + 1)
    model = max_mac_ll.Model(model.states.size, Nclass, clazz, pr)
    model.usePeq = Peq
    if notPeq:
        model.P0[:] = ModelToP0(file)
    model.useVar = True ##
    model.IeqFv = file.IeqFv
    model.Nchannel = file.channelCount
    model.iVoltage = data.iVoltage
    model.Gleak =Gleak
    model.Vleak =Vleak
    model.Vrev = file.vRev
    if model.IeqFv:
        for s in xrange(Nclass):
            model.amp[s] = file.classes.get(s, 'Cond')
            model.var[s] = file.classes.get(s, 'CondStd')**2
        model.amp0 = file.classes.get(0, 'Amp')
        model.var0 = file.classes.get(0, 'Std')**2
    else:
        for s in xrange(Nclass):
            model.amp[s] = file.classes.get(s, 'Amp')
            model.var[s] = file.classes.get(s, 'Std')**2
    K0, K1, L, V = ModelToRateMatrices(file, data.signal_names)
    model.K0[:] = K0[:]
    model.K1[:] = K1[:]
    model.L[:]  = L[:]
    model.V[:]  = V[:]
    model.setup_stim(data)
    A, B = file.constraints_kin.get_matrices(K0, K1, L, V)
    model.Ain[:A.shape[0],:A.shape[1]] = A
    model.Bin[:B.shape[0]] = B.flatten()
    model.Ncns = B.shape[0]
    model.setup_constraints()

    return model, file

def setup_files(data_path, model_path, delta):
    def pr(x):
        print x
    array = numpy.array
    
    data = setup_data(data_path, delta)
    model, model_file = setup_model(model_path, data)
    
    return model, data

if __name__ == '__main__':
    try:
        import datetime
        import optparse
        import sys
        import qubx.fast
        import qubx.model
        import qubx.model_constraints
        import qubx.data_types
        usage = "usage: %prog [options] file(s)"
        parser = optparse.OptionParser(usage)
        parser.add_option("-a", "--accel-level", dest="accel",
                          action="store",
                          default=2,
                          help="0: single-thread; 1: openmp; 2: opencl")
        parser.add_option("-r", "--resample-delta", dest="delta",
                          action="store", default=0.0,
                          help="if nonzero, resample-by delta")
        parser.add_option("-d", "--data", dest="data_path",
                          action="store", default="",
                          help="/path/to/data.qdf")
        parser.add_option("-m", "--model", dest="model_path",
                          action="store", default="",
                          help="/path/to/model.qmf")
        options, args = parser.parse_args()
        model, data = setup_files(options.data_path, options.model_path, options.delta)
        mac_ll(model, data, options.accel) # get it warmed up
        before = datetime.datetime.now()
        print mac_ll(model, data, options.accel),
        delta = datetime.datetime.now() - before
        print delta.seconds, (delta.microseconds / 1000), (delta.microseconds % 1000)
    except:
        traceback.print_exc()
        model, data = setup_test()
        print 'K0',model.K0
        print 'K1',model.K1
        print 'Nchannel',model.Nchannel
        print 'xamp',model.xamp
        print 'xvar',model.xvar
        print 'fPars',model.fPars
        print 'a'
        print mac_ll(model, data)
        print 'b'
        del data, model
        print 'c'