#!/usr/bin/env python3

# SynthMorph registration script using TensorFlow. To be used for affine and
# deformable registration and distributed with Docker.

import os
import sys
import time
import glob
import argparse
import numpy as np
import nibabel as nib
import tensorflow as tf

os.environ['VXM_BACKEND'] = 'tf'
os.environ['NEURITE_BACKEND'] = 'tf'

import neurite as ne
import voxelmorph as vxm


# Defaults.
in_shape = (256, 256, 256)

# Command-line arguments.
p = argparse.ArgumentParser()
p.add_argument('input')
p.add_argument('outdir')
args = p.parse_args()


def ori_to_ori(old, new='LIA', old_shape=None, zero_center=False):
    '''Construct matrix transforming coordinates from a voxel space with a new
    predominant anatomical axis orientation to an old orientation, by swapping
    and flipping axes. Operates in zero-based index space unless the space is
    to be zero-centered. The old shape must be specified if the old image is
    not a NiBabel object.'''
    def extract_ori(x):
        if isinstance(x, nib.filebasedimages.FileBasedImage):
            x = x.affine
        if isinstance(x, np.ndarray):
            return nib.orientations.io_orientation(x)
        if isinstance(x, str):
            return nib.orientations.axcodes2ornt(x)

    # Old shape.
    if zero_center:
        old_shape = (1, 1, 1)
    if old_shape is None:
        old_shape = old.shape

    # Transform from new to old index coordinates.
    old = extract_ori(old)
    new = extract_ori(new)
    new_to_old = nib.orientations.ornt_transform(old, new)
    return nib.orientations.inv_ornt_aff(new_to_old, old_shape)


def net_to_vox(im, out_shape=in_shape):
    '''Construct coordinate transform from isotropic 1-mm voxel space with
    gross LIA orentiation centered on the FOV - to the original image index
    space. The target space is a scaled and shifted voxel space, not world
    space.'''
    if isinstance(im, str):
        im = nib.load(im)

    # Gross LIA to predominant anatomical orientation of input image.
    assert isinstance(im, nib.filebasedimages.FileBasedImage) 
    lia_to_ori = ori_to_ori(im, new='LIA', old_shape=out_shape)

    # Scaling from millimeter to input voxels.
    vox_size = np.sqrt(np.sum(im.affine[:-1, :-1] ** 2, axis=0))
    scale = np.diag((*1 / vox_size, 1))

    # Shift from cen
    shift = np.eye(4)
    shift[:-1, -1] = 0.5 * (im.shape - out_shape / vox_size)

    # Total transform.
    return shift @ scale @ lia_to_ori


def transform(im, trans, shape=in_shape, normalize=False):
    '''Apply transformation matrix or field operating in zero-based index space
    to an image.'''
    if isinstance(im, nib.filebasedimages.FileBasedImage):
        im = im.get_fdata(dtype=np.float32)

    # Add singleton feature dimension if needed.
    if tf.rank(im) == 3:
        im = im[..., tf.newaxis]

    # Remove last row of matrix transforms.
    if tf.rank(trans) == 2 and trans.shape[0] == trans.shape[1]:
        trans = trans[:-1, :]

    out = vxm.utils.transform(
        im, trans, fill_value=0, shift_center=False, shape=shape,
    )

    if normalize:
        out -= tf.reduce_min(out)
        out /= tf.reduce_max(out)
    return out[tf.newaxis, ...]


def vm_key_weighted(
    in_shape=None,
    in_model=None,
    num_key=64,
    enc_nf=[256] * 4,
    dec_nf=[256] * 0,
    add_nf=[256] * 4,
    half_res=True,
    rigid=False,
):
    '''Find landmarks in images using a single-image detector and fit affine
    transform. Take mean as translation instead of fitting, and fit the
    transform in a centered frame.'''
    # Inputs.
    if in_model is None:
        source = tf.keras.Input(shape=(*in_shape, 1))
        target = tf.keras.Input(shape=(*in_shape, 1))
        in_model = tf.keras.Model(*[(source, target)] * 2)
    source, target = in_model.outputs[:2]

    in_shape = np.asarray(source.shape[1:-1])
    num_dim = len(in_shape)
    assert num_dim in (2, 3), 'only 2D and 3D supported'

    # Layers.
    down = getattr(tf.keras.layers, f'MaxPool{num_dim}D')()
    up = getattr(tf.keras.layers, f'UpSampling{num_dim}D')()
    act = tf.keras.layers.LeakyReLU(0.2)
    conv = getattr(tf.keras.layers, f'Conv{num_dim}D')
    prop = dict(kernel_size=3, padding='same')

    # Internal U-Net.
    inp = tf.keras.Input(shape=(*in_shape, 1))
    x = down(inp) if half_res else inp

    # Encoder.
    enc = []
    for n in enc_nf:
        x = conv(n, **prop)(x)
        x = act(x)
        enc.append(x)
        x = down(x)

    # Decoder.
    for n in dec_nf:
        x = conv(n, **prop)(x)
        x = act(x)
        x = tf.keras.layers.concatenate([up(x), enc.pop()])

    # Additional convolutions.
    for n in add_nf:
        x = conv(n, **prop)(x)
        x = act(x)

    # Features.
    x = conv(num_key, activation='relu', **prop)(x)
    net = tf.keras.Model(inp, outputs=x)
    key_1 = net(source)
    key_2 = net(target)

    # Barycenters.
    prop = dict(axes=range(1, num_dim + 1), normalize=True, shift_center=True)
    cen_1 = ne.utils.barycenter(key_1, **prop) * in_shape
    cen_2 = ne.utils.barycenter(key_2, **prop) * in_shape
    
    # Weights.
    axes = range(1, num_dim + 1)
    pow_1 = tf.reduce_sum(key_1, axis=axes)
    pow_2 = tf.reduce_sum(key_2, axis=axes)
    pow_1 /= tf.reduce_sum(pow_1, axis=-1, keepdims=True)
    pow_2 /= tf.reduce_sum(pow_2, axis=-1, keepdims=True)
    weights = pow_1 * pow_2

    # Least squares.
    out = vxm.utils.fit_affine(cen_1, cen_2, weights=weights)
    if rigid:
        out = vxm.utils.affine_matrix_to_params(out)
        out = out[:, :num_dim * (num_dim + 1) // 2]
        out = vxm.layers.ParamsToAffineMatrix(ndims=num_dim)(out)

    return tf.keras.Model(in_model.inputs, out)


# Setup.
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

# Model.
cpath = os.path.dirname(os.path.realpath(__file__))
modelfile = f'{cpath}/affine.h5'
model = vm_key_weighted(in_shape)
model.load_weights(modelfile)

os.makedirs(args.outdir, exist_ok=True)


# Affine align
for subj_path in glob.glob(os.path.join(args.input, "BraTSReg*")):

    subj = os.path.basename(subj_path)
    print(f"Performing alignment on {subj}")

    # Affine align
    sourcefile = glob.glob(f'{subj_path}/{subj}_00_????_t1.nii.gz')[0]
    targetfile = glob.glob(f'{subj_path}/{subj}_01_????_t1.nii.gz')[0]

    # Input data.
    mov = nib.load(sourcefile)
    fix = nib.load(targetfile)

    # Coordinate transforms. We will need these to take the images from their
    # native voxel spaces to network space. Voxel and network spaces are different
    # for each image. Network space is an isotropic 1-mm space centered on the
    # original image. Its axes are aligned with the original voxel data but flipped
    # and swapped to gross LIA orientation, which the network will expect.
    net_to_mov = net_to_vox(mov)
    net_to_fix = net_to_vox(fix)
    mov_to_net = np.linalg.inv(net_to_mov)
    fix_to_net = np.linalg.inv(net_to_fix)

    # Transforms from and to world space (RAS). There is only one world.
    mov_to_ras = mov.affine
    fix_to_ras = fix.affine
    ras_to_mov = np.linalg.inv(mov_to_ras)
    ras_to_fix = np.linalg.inv(fix_to_ras)

    # Transforms between zero-centered and zero-based voxel spaces.
    ind_to_cen = np.eye(4)
    ind_to_cen[:-1, -1] = -0.5 * (np.asarray(in_shape) - 1)
    cen_to_ind = np.eye(4)
    cen_to_ind[:-1, -1] = +0.5 * (np.asarray(in_shape) - 1)

    # Take the input images to network space. When saving the moving image with the
    # correct voxel-to-RAS matrix after incorporating an initial linear transform,
    # an image viewer taking this matrix into account will show an unchanged image.
    # However, the network only sees the voxel data, which have been moved.
    inputs = (
        transform(mov, net_to_mov, shape=in_shape, normalize=True),
        transform(fix, net_to_fix, shape=in_shape, normalize=True),
    )
    trans = model(inputs)

    # Add the last row to create a full matrix. Convert from zero-centered to
    # zero-based indices. Then compute the transform from native fixed to native
    # moving voxel spaces. Also compute a transform operating in RAS.
    trans = np.concatenate((np.squeeze(trans), np.reshape((0, 0, 0, 1), newshape=(1, -1))))
    trans = cen_to_ind @ trans @ ind_to_cen
    trans_vox = net_to_mov @ trans @ fix_to_net
    trans_ras = mov_to_ras @ trans_vox @ ras_to_fix

    # Output transforms operating in RAS.
    np.savetxt(fname=f'{args.outdir}/{subj}_affine.txt', X=trans_ras, fmt='%.8f %.8f %.8f %.8f')
