# Copyright 2014 Open Connectome Project (http://openconnecto.me)
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#     http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys
import tempfile
import numpy as np
import zlib
import cStringIO
from PIL import Image
import pylibmc
import time

import restargs
import ocpcadb
import ocpcaproj
import ocpcarest
import django
import re

from ocpca_cy import recolor_cy


class SimpleCatmaid:
  """Prefetch CATMAID tiles into MocpcacheDB"""

  def __init__(self):
    """Bind the mocpcache"""

    self.proj = None
    # make the mocpcache connection
    self.mc = pylibmc.Client(["127.0.0.1"], binary=True,behaviors={"tcp_nodelay":True,"ketama": True})

  def __del__(self):
    pass


  def buildKey (self,res,xtile,ytile,zslice):
    return 'simple/{}/{}/{}/{}/{}/{}'.format(self.token,self.tilesz,res,xtile,ytile,zslice)


  def cacheMiss ( self, resolution, xtile, ytile, zslice ):
    """On a miss. Cutout, return the image and load the cache in a background thread"""

    # make sure that the tile size is aligned with the cubedim
    if self.tilesz % self.proj.datasetcfg.cubedim[resolution][0] != 0 or self.tilesz % self.proj.datasetcfg.cubedim[resolution][1]:
      raise("Illegal tile size.  Not aligned")

    # figure out the cutout (limit to max image size)
    xstart = xtile*self.tilesz
    ystart = ytile*self.tilesz
    xend = min ((xtile+1)*self.tilesz,self.proj.datasetcfg.imagesz[resolution][0])
    yend = min ((ytile+1)*self.tilesz,self.proj.datasetcfg.imagesz[resolution][1])

    # get an xy image slice
    imageargs = '{}/{},{}/{},{}/{}/'.format(resolution,xstart,xend,ystart,yend,zslice) 
    cb = ocpcarest.xySlice ( imageargs, self.proj, self.db )
    if cb.data.shape != (1,self.tilesz,self.tilesz):
      tiledata = np.zeros((self.tilesz,self.tilesz), cb.data.dtype )
      tiledata[0:((yend-1)%self.tilesz+1),0:((xend-1)%self.tilesz+1)] = cb.data[0,:,:]
    else:
      tiledata = cb.data

    # need to make polymorphic for different image types     
    if tiledata.dtype == np.uint8:
      outimage = Image.frombuffer ( 'L', (self.tilesz,self.tilesz), tiledata, 'raw', 'L', 0, 1 ) 
    elif tiledata.dtype == np.uint32:
      tiledata = tiledata.reshape([self.tilesz,self.tilesz])
      recolor_cy (tiledata, tiledata)
      outimage = Image.frombuffer ( 'RGBA', [self.tilesz,self.tilesz], tiledata, 'raw', 'RGBA', 0, 1 )
    else:
      #added by PJM to support probmap type.
      tiledata = np.uint8(tiledata*256)
      outimage = Image.frombuffer ( 'L', (self.tilesz,self.tilesz), tiledata, 'raw', 'L', 0, 1 ) 
      #assert 0 
      # need to fix here and add falsecolor args

    return outimage


  def getTile ( self, webargs ):
    """Either fetch the file from mocpcache or get a mcfc image"""

    # parse the web args
    self.token, tileszstr, resstr, xtilestr, ytilestr, zslicestr, rest = webargs.split('/',6)
    [ self.db, self.proj, projdb ] = ocpcarest.loadDBProj ( self.token )

    # convert args to ints
    xtile = int(xtilestr)
    ytile = int(ytilestr)
    res = int(resstr)
    # xyslice will modify zslice to the offset
    zslice = int(zslicestr)
    self.tilesz = int(tileszstr)

    # mocpcache key
    mckey = self.buildKey(res,xtile,ytile,zslice)

    # do something to sanitize the webargs??
    # if tile is in mocpcache, return it
    tile = self.mc.get(mckey)
    if tile == None:
      img=self.cacheMiss(res,xtile,ytile,zslice)
      fobj = cStringIO.StringIO ( )
      img.save ( fobj, "PNG" )
      self.mc.set(mckey,fobj.getvalue())
    else:
      fobj = cStringIO.StringIO(tile)

    fobj.seek(0)
    return fobj

@@ -2410,16 +2410,59 @@
 b.data%0A%0A
+    cb.catmaidSlice( )%0A%0A    return cb.data%0A
     # ne
@@ -2515,24 +2515,25 @@
 es     %0A    
+#
 if tiledata.
@@ -2551,24 +2551,25 @@
 .uint8:%0A    
+#
   outimage =
@@ -2654,16 +2654,17 @@
  ) %0A    
+#
 elif til
@@ -2689,24 +2689,25 @@
 uint32:%0A    
+#
   tiledata =
@@ -2755,16 +2755,17 @@
 z%5D)%0A    
+#
   recolo
@@ -2790,24 +2790,25 @@
 ledata)%0A    
+#
   outimage =
@@ -2894,24 +2894,25 @@
  0, 1 )%0A    
+#
 else:%0A      
@@ -2905,24 +2905,25 @@
  #else:%0A    
+#
   #added by 
@@ -2951,24 +2951,25 @@
 p type.%0A    
+#
   tiledata =
@@ -2996,16 +2996,17 @@
 56)%0A    
+#
   outima
@@ -3162,29 +3162,8 @@
 gs%0A%0A
-    return outimage%0A%0A
 %0A  d

Support for other datatypes in catmaid