"""
"""
from __future__ import absolute_import, print_function

# C:\home\ej\wrk\scipy\weave\examples>python vq.py
# vq with 1000 observation, 10 features and 30 codes fo 100 iterations
#  speed in python: 0.150119999647
# [25 29] [ 2.49147266  3.83021032]
#  speed in standard c: 0.00710999965668
# [25 29] [ 2.49147266  3.83021032]
#  speed up: 21.11
#  speed inline/blitz: 0.0186300003529
# [25 29] [ 2.49147272  3.83021021]
#  speed up: 8.06
#  speed inline/blitz2: 0.00461000084877
# [25 29] [ 2.49147272  3.83021021]
#  speed up: 32.56

from numpy import *
import sys
sys.path.insert(0,'..')
import scipy.weave.inline_tools as inline_tools
import scipy.weave.converters as converters
blitz_type_converters = converters.blitz
import scipy.weave.c_spec as c_spec

def vq(obs,code_book):
    # make sure we're looking at arrays.
    obs = asarray(obs)
    code_book = asarray(code_book)
    # check for 2d arrays and compatible sizes.
    obs_sh = shape(obs)
    code_book_sh = shape(code_book)
    assert(len(obs_sh) == 2 and len(code_book_sh) == 2)
    assert(obs_sh[1] == code_book_sh[1])
    type = c_spec.num_to_c_types[obs.typecode()]
    # band aid for now.
    ar_type = 'PyArray_FLOAT'
    code =  """
            #line 37 "vq.py"
            // Use tensor notation.
            blitz::Array<%(type)s,2> dist_sq(Ncode_book[0],Nobs[0]);
                blitz::firstIndex i;
            blitz::secondIndex j;
            blitz::thirdIndex k;
            dist_sq = sum(pow2(obs(j,k) - code_book(i,k)),k);
            // Surely there is a better way to do this...
            PyArrayObject* py_code = (PyArrayObject*) PyArray_FromDims(1,&Nobs[0],PyArray_LONG);
                blitz::Array<int,1> code((int*)(py_code->data),
                                     blitz::shape(Nobs[0]), blitz::neverDeleteData);
                code = minIndex(dist_sq(j,i),j);

                PyArrayObject* py_min_dist = (PyArrayObject*) PyArray_FromDims(1,&Nobs[0],PyArray_FLOAT);
                blitz::Array<float,1> min_dist((float*)(py_min_dist->data),
                                               blitz::shape(Nobs[0]), blitz::neverDeleteData);
                min_dist = sqrt(min(dist_sq(j,i),j));
                py::tuple results(2);
                results[0] = py_code;
                results[1] = py_min_dist;
                return_val = results;
            """ % locals()
    code, distortion = inline_tools.inline(code,['obs','code_book'],
                                           type_converters = blitz_type_converters,
                                           compiler = 'gcc',
                                           verbose = 1)
    return code, distortion

def vq2(obs,code_book):
    """ doesn't use blitz (except in conversion)
        ALSO DOES NOT HANDLE STRIDED ARRAYS CORRECTLY
    """
    # make sure we're looking at arrays.
    obs = asarray(obs)
    code_book = asarray(code_book)
    # check for 2d arrays and compatible sizes.
    obs_sh = shape(obs)
    code_book_sh = shape(code_book)
    assert(len(obs_sh) == 2 and len(code_book_sh) == 2)
    assert(obs_sh[1] == code_book_sh[1])
    assert(obs.typecode() == code_book.typecode())
    type = c_spec.num_to_c_types[obs.typecode()]
    # band aid for now.
    ar_type = 'PyArray_FLOAT'
    code =  """
            #line 83 "vq.py"
            // THIS DOES NOT HANDLE STRIDED ARRAYS CORRECTLY
            // Surely there is a better way to do this...
            PyArrayObject* py_code = (PyArrayObject*) PyArray_FromDims(1,&Nobs[0],PyArray_LONG);
                PyArrayObject* py_min_dist = (PyArrayObject*) PyArray_FromDims(1,&Nobs[0],PyArray_FLOAT);