# @author: msharrock
# version: 0.0.1

"""
Registration methods for DeepBleed 

"""
import ants

def rigid(fixed, moving):

        '''
        Rigid Registration with ANTs

        Params:
                - moving: ants image, image to move when registering
                - fixed: ants image, template image to register to

        Outputs: 
                - image: registered image
                - transforms: transformation affine matrix
        '''

        kwargs = {'-n': 'nearestNeighbor'}
        tx = ants.registration(fixed, moving, type_of_transform='Rigid', mask=None, grad_step=0.2, flow_sigma=3, total_sigma=0, 
                           aff_metric='mattes', aff_sampling=64, syn_metric ='mattes',**kwargs) 
                        
        image = tx['warpedmovout']
        transforms = tx['fwdtransforms']
        return image, transforms


def invert(fixed, moving, transforms):
        '''
        Inverse Transform with ANTs

        Params:
                - image: ants image, image to revert
                - invtransform: affine matrix to use for inverse transform
        Outputs: 
                - image: ants image, inverted
        '''
        image = ants.apply_transforms(fixed = fixed, moving = moving, transformlist = transforms, interpolator = 'genericLabel', whichtoinvert = [True])   

        return image