import sys
sys.path.append('../code/functions')
sys.path.append('/home/simpleElastix/build/SimpleITK-build/Wrapping/Python')

import pickle
import cv2
import time

import SimpleITK as sitk
import numpy as np
import matplotlib.pyplot as plt
import nibabel as nib

from cluster import Cluster
from tiffIO import loadTiff, unzipChannels
from connectLib import adaptiveThreshold
from hyperReg import get3DRigid, parseTransformFile, apply3DRigid, apply3DRigidVolume

movingImg = pickle.load(open('../code/functions/movingDebug.io', 'r'))
fixedImg = pickle.load(open('../code/functions/fixedDebug.io', 'r'))
movingLmk = pickle.load(open('../code/functions/movingLDebug.io', 'r'))
fixedLmk = pickle.load(open('../code/functions/fixedLDebug.io', 'r'))

testFixedLmk = fixedLmk[10:15]
testMovingLmk = movingLmk[10:15]
testFixedImg = fixedImg[10:15]
testMovingImg = movingImg[10:15]

plt.figure()
plt.imshow(movingImg[10], cmap='gray')
plt.title('Moving Image Post Connected Components')
plt.show()

plt.figure()
plt.imshow(movingImg[10], cmap='gray')
plt.title('Fixed Image Post Connected Components')
plt.show()

plt.figure()
plt.imshow(movingLmk[10], cmap='gray')
plt.title('Moving Image Landmarks')
plt.show()

plt.figure()
plt.imshow(fixedLmk[10], cmap='gray')
plt.title('Fixed Image Landmarks')
plt.show()

"""
Explanation: Load Data
In order to expediate the testing process, I added a debug flag to the pipeline method in our pipeline file which outputs the fixed and moving images prior to registration
End of explanation
"""


def toDiff(imgA, imgB):
    ret = np.empty((imgA.shape[0], imgA.shape[1], 3), dtype=np.uint8)
    for y in range(imgA.shape[0]):
        for x in range(imgA.shape[1]):
            
            if imgA[y][x] and not imgB[y][x]:
                ret[y][x][0] = 255
                ret[y][x][1] = 0
                ret[y][x][2] = 0
            elif not imgA[y][x] and imgB[y][x]:
                ret[y][x][0] = 0
                ret[y][x][1] = 255
                ret[y][x][2] = 0
            elif imgA[y][x] and imgB[y][x]:
                ret[y][x][0] = 255
                ret[y][x][1] = 0
                ret[y][x][2] = 255
            else:
                ret[y][x][0] = 255
                ret[y][x][1] = 255
                ret[y][x][2] = 255
            
    return ret

def visDiff(sliceA, sliceB):
    disp = toDiff(sliceA, sliceB)
    return disp

def visVolDiff(volumeA, volumeB):
    for i in range(volumeA.shape[0]):
        plt.figure()
        plt.title('Disperity at z=' + str(i))
        plt.imshow(visDiff(volumeA[i], volumeB[i]))
        plt.show()

"""
Explanation: Visualization Function
End of explanation
"""


def preproc(img):
    binImg = adaptiveThreshold(img, 5, 5)
    binImg*=255
    outImg = np.stack([cv2.erode(sub, None, 1) for sub in binImg])
    outImg = np.stack([cv2.dilate(sub, None, 2) for sub in outImg])

    return outImg

def register(landmarks1, landmarks2, additionalParams):
    SimpleElastix = sitk.SimpleElastix()

    SimpleElastix.LogToConsoleOn()

    img1 = nib.Nifti1Image(preproc(landmarks1), np.eye(4))
    nib.save(img1, 'fixed.nii')
    img2 = nib.Nifti1Image(preproc(landmarks2), np.eye(4))
    nib.save(img2, 'moving.nii')

    SimpleElastix.SetFixedImage(sitk.ReadImage('fixed.nii'))
    SimpleElastix.SetMovingImage(sitk.ReadImage('moving.nii'))

    pMap = sitk.GetDefaultParameterMap('rigid')
    for elem in additionalParams:
        pMap[elem[0]]=[elem[1]]
    SimpleElastix.SetParameterMap(pMap)
    SimpleElastix.Execute()
    
    t = SimpleElastix.GetTransformParameterMap()
    sitk.WriteParameterFile(t[0], 'transform.txt')
    imgFilter = sitk.SimpleTransformix()
    imgFilter.SetTransformParameterMap(t[0])
    imgFilter.PrintParameterMap()
    return imgFilter

start = time.time()
myTransform = register(fixedLmk, movingLmk,
                        [['MaximumNumberOfSamplingAttempts','200'],
                         ['MaximumNumberOfIterations', '700'],
                         ['Metric', 'AdvancedMeanSquares']]
                      )
end = time.time()

print end - start

params = parseTransformFile('transform.txt')

regVolume = apply3DRigidVolume(fixedLmk,
                               params['matrix'],
                               params['originZ'],
                               params['originY'],
                               params['originX'])

visVolDiff(preproc(fixedLmk[140:145]), preproc(movingLmk[140:145]))

visVolDiff(preproc(regVolume[140:145]), preproc(movingLmk[140:145]))

"""
Explanation: Registration Functions
End of explanation
"""


