from __future__ import division
import numpy as np
cimport numpy as np
cimport cython

DTYPE1 = np.float
ctypedef np.float_t DTYPE1_t
DTYPE2 = np.int
ctypedef np.int_t DTYPE2_t
DTYPE3 = np.float32
ctypedef np.float32_t DTYPE3_t

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
#####################################################################################################
# Caveat:
# If the feature you are looking for has over 100000 points, this will Segfault. 
# We can increase that number, but at a large decrease in efficiency (or perhaps this code needs to be
# designed better). If you are searching a large 2D grid for a value that occurs frequently, 
# then just use numpy.where.
def cython_where(np.ndarray[DTYPE3_t, ndim=2] data, oper, DTYPE1_t val, np.ndarray[DTYPE2_t, ndim=1] xind, np.ndarray[DTYPE2_t, ndim=1] yind):
  assert data.dtype == DTYPE3
  
  cdef int xmax = data.shape[0]
  cdef int ymax = data.shape[1]
  cdef unsigned int x, y
  cdef int count = 0
#  cdef np.ndarray[DTYPE2_t, ndim=1] xind = np.empty(100000,dtype=int)
#  cdef np.ndarray[DTYPE2_t, ndim=1] yind = np.empty(100000,dtype=int)
  if(oper == 'EQ' or oper == 'eq'):
    for x in xrange(xmax):
      for y in xrange(ymax):
        if(data[x,y] == val):
          xind[count] = x
          yind[count] = y
          count += 1
  elif(oper == 'GE' or oper == 'ge'):
    for x in xrange(xmax):
      for y in xrange(ymax):
        if(data[x,y] >= val):
          xind[count] = x
          yind[count] = y
          count += 1
  elif(oper == 'GT' or oper == 'gt'):
    for x in xrange(xmax):
      for y in xrange(ymax):
        if(data[x,y] > val):
          xind[count] = x
          yind[count] = y
          count += 1
  elif(oper == 'LE' or oper == 'le'):
    for x in xrange(xmax):
      for y in xrange(ymax):
        if(data[x,y] <= val):
          xind[count] = x
          yind[count] = y
          count += 1
  else: # LT
    for x in xrange(xmax):
      for y in xrange(ymax):
        if(data[x,y] < val):
          xind[count] = x
          yind[count] = y
          count += 1
        

  return tuple([xind[0:count],yind[0:count]]),count
#####################################################################################################
def cython_where_EQ_1D(np.ndarray[DTYPE1_t, ndim=1] data, DTYPE1_t val):
  assert data.dtype == DTYPE1

  cdef int xmax = data.shape[0]
  cdef unsigned int x
  cdef int count = 0
  cdef np.ndarray[DTYPE2_t, ndim=1] xind = np.zeros(100000,dtype=int)
  for x in xrange(xmax):
    if(data[x] == val):
      xind[count] = x
      count += 1


  return xind[0:count],count

#####################################################################################################