#!/usr/bin/env python
# coding: utf-8

# #  Multiscale Manager
# 
# Here we will briefly describe some of the features of the multiscale manager, using some very simple examples. 
# 
# Multiscale processing involves estimating using lower resolution data (for speed and stability), and then estimate on higher and higher resolutions.  Therefore, for multiscale processing, we downsample our data at each scale level (with the exception of the final, full-res level). We also upsample all our estimates everytime we switch levels.
# 
# In image registration, for example, our data is two (or more) full-res images, and our estimates are velocity fields. We do the following for each scale level:
# 
# 1. For the first scale (smallest grid) 
#     * downsample full-res data to grid
#     * initialize the estimates
# 2. For all intermediate scales:
#     * downsample full-res data to grid
#     * upsample the previous estimates
# 3. For the last scale (full resolution):
#     * use the full-res data
#     * upsample the previous estimates
#     
# In PyCA, the MultiscaleManager keeps track of the scales and associated grid sizes, and the multiscale resampler allows functionality to upsample/downsample images, v-fields, and h-fields.
# 
# Since we will eventually need memory to support these full-res operations, we define all the Image3Ds and Field3Ds at full-res, and then we just use subsets of their memory at lower scales. This way we only need to allocate memory once.
# 
# For this dummy example, we assume that we have 1 image as data, and we are trying to estimate a v-field and an h-field. 

# In[1]:


import PyCA.Core as ca

origGrid = ca.GridInfo(ca.Vec3Di(100, 100, 1))
mType = ca.MEM_DEVICE

#Data image
I0Data = ca.Image3D(origGrid, mType) #Full res Image
ca.SetMem(I0Data, 2.0)

#Downsampled data - allocate at full-res grid
I0 = ca.Image3D(origGrid, mType) #(possibly) downsampled version of I0Orig

#Estimates - allocate at full-res grid
v = ca.Field3D(origGrid, mType) # v-field to estimate
h = ca.Field3D(origGrid, mType) # h-field to estimate
Iest = ca.Image3D(origGrid, mType) # Image3D to estimate

# temporary *3Ds - allocate at full-res grid
Itemp0 = ca.Image3D(origGrid, mType) #temporary image
Itemp1 = ca.Image3D(origGrid, mType)


# Next, we define the number of scales we want (we are choosing a downsampling of 4:1, 2:1, and 1:1 (full res).
# 
# We now initialize the scale manager and the resampler.

# In[2]:


scales = [4, 2, 1]
scaleManager = ca.MultiscaleManager(origGrid)
for s in scales:
    scaleManager.addScaleLevel(s)
    
if mType == ca.MEM_HOST:
    resampler = ca.MultiscaleResamplerGaussCPU(origGrid)
else:
    resampler = ca.MultiscaleResamplerGaussGPU(origGrid)


# We now define what happens everytime we change scales. It is typically easiest to write this as a function that takes the scale level. The scale manager will let us check if we are on the first or the last scale. Note: the scale level is an index (in this case, scale level only accepts 0, 1, 2).

# In[3]:


def setScale(scale):
    scaleManager.set(scale)
    resampler.setScaleLevel(scaleManager)
    curGrid = scaleManager.getCurGrid()
    curGrid.spacing().z = 1.0 #since we're only doing 2D
    print 'Inside setScale(). Current grid is ', curGrid
    
    # set grid for all other Image3D/Field3D objects
    Itemp0.setGrid(curGrid)
    Itemp1.setGrid(curGrid)
    
    
    if scaleManager.isFirstScale():
        print 'Inside setScale: **First scale**'
    if scaleManager.isLastScale():
        print 'Inside setScale: **First scale**'   
    
    # downsample data (if not last scale)
    I0.setGrid(curGrid)
    if scaleManager.isLastScale():
        ca.Copy(I0, I0Data)
    else:
        resampler.downsampleImage(I0, I0Data)
        
    # initialize estimates (for first scale), otherwise upsample
    if scaleManager.isFirstScale():
        h.setGrid(curGrid)
        v.setGrid(curGrid)
        Iest.setGrid(curGrid)
        ca.SetToIdentity(h)
        ca.SetToZero(v)
        ca.SetMem(Iest, 0.0)
    else:
        resampler.updateHField(h)
        resampler.updateVField(v)
        resampler.upsampleImage(Iest)    


# One thing to note:  If you upsample an image, it will automatically change the grid for you, so don't change the grid before you upsample!.
# 
# Now we show some simple code that goes through the scales and updates the images/fields.

# In[4]:


setScale(0)
print "In scale 0, range of v is", ca.MinMax(v)
print "In scale 0, range of Iest is", ca.MinMax(Iest)
v += 1.0
Iest += 1.0
setScale(1)
print "In scale 1, range of v is", ca.MinMax(v)
print "In scale 1, range of Iest is", ca.MinMax(Iest)
v += 1.0
Iest += 1.0
setScale(2)
print "In final scale, range of v is", ca.MinMax(v)
print "In final scale, range of Iest is", ca.MinMax(v)


# There are a number of other multiscale functions.  For completeness, I will list them here. They are pretty self explanatory.

# In[5]:


setScale(0)
setScale(1)
print '(Current) number of voxels:', scaleManager.curNVox()
print '(Current) voxel volume:', scaleManager.curVoxelVol()
print '(Original) number of voxels:', scaleManager.origNVox()
print '(Original) voxel volume:', scaleManager.origVoxelVol()

print 'Current scale factor (new spacing/previous spacing):', scaleManager.getCurFactor()
print 'Current scale level (int):', scaleManager.getCurLevel()
print 'Current grid:', scaleManager.getCurGrid(), scaleManager.getCurSize(), scaleManager.getCurSpace()
print 'Original grid:',  scaleManager.getOrgGrid(), scaleManager.getOrgSize(), scaleManager.getOrgSpace()
print 'booleans:', scaleManager.isFirstScale(), scaleManager.isIdentityScale(), scaleManager.isLastScale()