#!/usr/bin/env python3

import csv 

import os
import sys 

#from nifti import *;
#import nifti.clib as ncl 

import nibabel

import numpy
from numpy import linalg as LA

import scipy.ndimage
import math

import getopt

import scipy.io
import scipy.ndimage

import tempfile
import shutil
import subprocess
import FLIRT

def usage():
    print("DWIEddyCorrectMI [-p] [-m <output mat file>] <input file> <output file> <reference index>")
    print()
    print("Performs alignment of each volume in <input file> to the image in the input file specified by <reference index>; the first volume is 0.")
    print()
    print("Options:")
    print("\t-m <output mat file>: contains all of the transformation matrices concatenated")
    print("\t-p use GNU parallel to perform the registrations in parallel, uses GNU parallel")

def gluInvertMatrix(m):
    TestInv = numpy.zeros([4, 4])

    TestInv[0, 0] = m[1, 1]  * m[2, 2] * m[3, 3] - \
    m[1, 1]  * m[3, 2] * m[2, 3] -  \
    m[1, 2]  * m[2, 1]  * m[3, 3] +  \
    m[1, 2]  * m[3, 1]  * m[2, 3] + \
    m[1, 3] * m[2, 1]  * m[3, 2] -  \
    m[1, 3] * m[3, 1]  * m[2, 2]

    TestInv[0, 1] = -m[0, 1]  * m[2, 2] * m[3, 3] +  \
    m[0, 1]  * m[3, 2] * m[2, 3] +  \
    m[0, 2]  * m[2, 1]  * m[3, 3] -  \
    m[0, 2]  * m[3, 1]  * m[2, 3] -  \
    m[0, 3] * m[2, 1]  * m[3, 2] +  \
    m[0, 3] * m[3, 1]  * m[2, 2]

    TestInv[0, 2] = m[0, 1]  * m[1, 2] * m[3, 3] -  \
    m[0, 1]  * m[3, 2] * m[1, 3] -  \
    m[0, 2]  * m[1, 1] * m[3, 3] +  \
    m[0, 2]  * m[3, 1] * m[1, 3] +  \
    m[0, 3] * m[1, 1] * m[3, 2] -  \
    m[0, 3] * m[3, 1] * m[1, 2]

    TestInv[0, 3] = -m[0, 1]  * m[1, 2] * m[2, 3] +  \
    m[0, 1]  * m[2, 2] * m[1, 3] +  \
    m[0, 2]  * m[1, 1] * m[2, 3] -  \
    m[0, 2]  * m[2, 1] * m[1, 3] -  \
    m[0, 3] * m[1, 1] * m[2, 2] +  \
    m[0, 3] * m[2, 1] * m[1, 2]

    TestInv[1, 0] = -m[1, 0]  * m[2, 2] * m[3, 3] +  \
    m[1, 0]  * m[3, 2] * m[2, 3] +  \
    m[1, 2]  * m[2, 0] * m[3, 3] -  \
    m[1, 2]  * m[3, 0] * m[2, 3] -  \
    m[1, 3] * m[2, 0] * m[3, 2] +  \
    m[1, 3] * m[3, 0] * m[2, 2]

    TestInv[1, 1] = m[0, 0]  * m[2, 2] * m[3, 3] -  \
    m[0, 0]  * m[3, 2] * m[2, 3] -  \
    m[0, 2]  * m[2, 0] * m[3, 3] +  \
    m[0, 2]  * m[3, 0] * m[2, 3] +  \
    m[0, 3] * m[2, 0] * m[3, 2] -  \
    m[0, 3] * m[3, 0] * m[2, 2]

    TestInv[1, 2] = -m[0, 0]  * m[1, 2] * m[3, 3] +  \
    m[0, 0]  * m[3, 2] * m[1, 3] +  \
    m[0, 2]  * m[1, 0] * m[3, 3] -  \
    m[0, 2]  * m[3, 0] * m[1, 3] -  \
    m[0, 3] * m[1, 0] * m[3, 2] +  \
    m[0, 3] * m[3, 0] * m[1, 2]

    TestInv[1, 3] = m[0, 0]  * m[1, 2] * m[2, 3] -  \
    m[0, 0]  * m[2, 2] * m[1, 3] -  \
    m[0, 2]  * m[1, 0] * m[2, 3] +  \
    m[0, 2]  * m[2, 0] * m[1, 3] +  \
    m[0, 3] * m[1, 0] * m[2, 2] -  \
    m[0, 3] * m[2, 0] * m[1, 2]

    TestInv[2, 0] = m[1, 0]  * m[2, 1] * m[3, 3] -  \
    m[1, 0]  * m[3, 1] * m[2, 3] -  \
    m[1, 1]  * m[2, 0] * m[3, 3] +  \
    m[1, 1]  * m[3, 0] * m[2, 3] +  \
    m[1, 3] * m[2, 0] * m[3, 1] -  \
    m[1, 3] * m[3, 0] * m[2, 1]

    TestInv[2, 1] = -m[0, 0]  * m[2, 1] * m[3, 3] +  \
    m[0, 0]  * m[3, 1] * m[2, 3] +  \
    m[0, 1]  * m[2, 0] * m[3, 3] -  \
    m[0, 1]  * m[3, 0] * m[2, 3] -  \
    m[0, 3] * m[2, 0] * m[3, 1] +  \
    m[0, 3] * m[3, 0] * m[2, 1]

    TestInv[2, 2] = m[0, 0]  * m[1, 1] * m[3, 3] -  \
    m[0, 0]  * m[3, 1] * m[1, 3] -  \
    m[0, 1]  * m[1, 0] * m[3, 3] +  \
    m[0, 1]  * m[3, 0] * m[1, 3] +  \
    m[0, 3] * m[1, 0] * m[3, 1] -  \
    m[0, 3] * m[3, 0] * m[1, 1]

    TestInv[2, 3] = -m[0, 0]  * m[1, 1] * m[2, 3] +  \
    m[0, 0]  * m[2, 1] * m[1, 3] +  \
    m[0, 1]  * m[1, 0] * m[2, 3] -  \
    m[0, 1]  * m[2, 0] * m[1, 3] -  \
    m[0, 3] * m[1, 0] * m[2, 1] +  \
    m[0, 3] * m[2, 0] * m[1, 1]

    TestInv[3, 0] = -m[1, 0] * m[2, 1] * m[3, 2] +  \
    m[1, 0] * m[3, 1] * m[2, 2] +  \
    m[1, 1] * m[2, 0] * m[3, 2] -  \
    m[1, 1] * m[3, 0] * m[2, 2] -  \
    m[1, 2] * m[2, 0] * m[3, 1] +  \
    m[1, 2] * m[3, 0] * m[2, 1]

    TestInv[3, 1] = m[0, 0] * m[2, 1] * m[3, 2] -  \
    m[0, 0] * m[3, 1] * m[2, 2] -  \
    m[0, 1] * m[2, 0] * m[3, 2] +  \
    m[0, 1] * m[3, 0] * m[2, 2] +  \
    m[0, 2] * m[2, 0] * m[3, 1] -  \
    m[0, 2] * m[3, 0] * m[2, 1]

    TestInv[3, 2] = -m[0, 0] * m[1, 1] * m[3, 2] +  \
    m[0, 0] * m[3, 1] * m[1, 2] +  \
    m[0, 1] * m[1, 0] * m[3, 2] -  \
    m[0, 1] * m[3, 0] * m[1, 2] -  \
    m[0, 2] * m[1, 0] * m[3, 1] +  \
    m[0, 2] * m[3, 0] * m[1, 1]

    TestInv[3, 3] = m[0, 0] * m[1, 1] * m[2, 2] -  \
    m[0, 0] * m[2, 1] * m[1, 2] -  \
    m[0, 1] * m[1, 0] * m[2, 2] +  \
    m[0, 1] * m[2, 0] * m[1, 2] +  \
    m[0, 2] * m[1, 0] * m[2, 1] -  \
    m[0, 2] * m[2, 0] * m[1, 1]

    TestDet = m[0, 0] * TestInv[0, 0] + m[1, 0] * TestInv[0, 1] + m[2, 0] * TestInv[0, 2] + m[3, 0] * TestInv[0, 3]

    if (TestDet == 0):
        OutInv = numpy.array([4, 4])
        OutInv.fill(numpy.nan)
    else:
        TestDet = 1.0 / TestDet
        OutInv = numpy.multiply(TestInv, TestDet)
    return OutInv

#def main():
        
opts, args = getopt.getopt(sys.argv[1:], "m:p", [])
    
if len(args) < 2:
    usage()
    quit()

OutputMATFile = None
parallelMode = False

for o, a in opts:
    if o == '-m':
        OutputMATFile = a
    if o == '-p':
        parallelMode = True

DWIImageName = args[0]
#GradientTableFile = args[1]
OutputFile = args[1]
try:
    ReferenceIDX = int(args[2])
except e:
    print("The reference index was not an integer")
    quit()

#NI = NiftiImage(DWIImageName)

#NumTotalVolumes = NI.data.shape[0]
#NumSlices = NI.data.shape[1]
#NumVoxels = NI.data.shape[2] * NI.data.shape[3]

NIB = nibabel.load(DWIImageName)

#print "NI--------------------"
#print NI
#print "NIB--------------------"
#print NIB
#print NIB.shape

NumTotalVolumes = NIB.shape[3]
#NumSlices = NIB.shape[2]
#NumVoxels = NIB.shape[0] * NIB.shape[1]

del NIB

#quit()
if ReferenceIDX < 0 or ReferenceIDX >= NumTotalVolumes:
    print("Reference index is out of bounds")
    exit()

#GradientTableFileFID = open(GradientTableFile, 'r')
#GradientTableFileReader = csv.reader(GradientTableFileFID, delimiter='\t')

#GradientTable = numpy.zeros([NumTotalVolumes, 3])

#NonZeroBIDX = list()
#NonBZeroVolumes = list()
#BZeroVolumes = list()

#ImageIDX = 0
#for row in GradientTableFileReader:
#    if not row[3] == '0':
#        NonZeroBIDX.append(ImageIDX)
#        NonBZeroVolumes.append(ImageIDX)
#    else:
#        BZeroVolumes.append(ImageIDX)
#
#    GradientTable[ImageIDX, :] = [row[0], row[1], row[2]]
#    ImageIDX = ImageIDX + 1

#GradientTableFileFID.close()

#print GradientTable

Transforms = numpy.zeros([NumTotalVolumes, 4, 4])
InvTransforms = numpy.zeros([NumTotalVolumes, 4, 4])

HeadTranslationsPlain = numpy.zeros([NumTotalVolumes, 3])
HeadScalesPlain = numpy.zeros([NumTotalVolumes, 3])
HeadAnglesPlain = numpy.zeros([NumTotalVolumes, 3])
HeadTranslationsBET = numpy.zeros([NumTotalVolumes, 3])
HeadScalesBET = numpy.zeros([NumTotalVolumes, 3])
HeadAnglesBET = numpy.zeros([NumTotalVolumes, 3])

#TempDirPlain = tempfile.mkdtemp()
#TempDirBET = tempfile.mkdtemp()

TempDir = tempfile.mkdtemp()#os.path.join(OutputDirRoot, 'plain')

subprocess.call(['fslsplit', DWIImageName, os.path.join(TempDir, "tmp")])

#subprocess.call(['fslroi', DWIImageName, os.path.join(TempDir, "tmp0000"), '0', '1'])

ReferenceVolume = os.path.join(TempDir, "tmp0000")
subprocess.call(['bet', ReferenceVolume, ReferenceVolume + "_brain", '-f', '0.3', '-m'])
#print(ReferenceVolume + "_brain")

#subprocess.call(['fslmaths', ReferenceVolume + "_brain_mask", '-eroF', '-eroF', '-s', '4', ReferenceVolume + "_brain_mask_weighted", '-odt', 'float'])
#subprocess.call(['fslmaths', ReferenceVolume + "_brain_mask_weighted", '-eroF', '-eroF', '-s', '8', '-dilF', '-dilF', ReferenceVolume + "_brain_mask_weighted_fallback", '-odt', 'float'])
#print("erode")
#subprocess.call(['fslmaths', ReferenceVolume + "_brain_mask_weighted", '-eroF', ReferenceVolume + "_brain_mask_weighted_fallback", '-odt', 'float'])
#print("smooth")
#subprocess.call(['fslmaths', ReferenceVolume + "_brain_mask_weighted_fallback", '-s', '8', ReferenceVolume + "_brain_mask_weighted_fallback", '-odt', 'float'])

#print("dilate")
#subprocess.call(['fslmaths', ReferenceVolume + "_brain_mask_weighted_fallback", '-dilF', ReferenceVolume + "_brain_mask_weighted_fallback", '-odt', 'float'])
#print("done")
BrainMaskNII = nibabel.load(ReferenceVolume + "_brain_mask.nii.gz")
BrainMaskIMG = BrainMaskNII.get_fdata() > 0
BrainMaskWeightedIMG = scipy.ndimage.binary_erosion(BrainMaskIMG, iterations = 2, structure = numpy.ones((3, 3, 3)))
BrainMaskWeightedIMG = scipy.ndimage.gaussian_filter(numpy.single(BrainMaskWeightedIMG), sigma = [3.2 / x for x in BrainMaskNII.header.get_zooms()])
BrainMaskWeightedIMGFallback = scipy.ndimage.grey_erosion(BrainMaskWeightedIMG, footprint = numpy.ones((3, 3, 3)))
BrainMaskWeightedIMGFallback = scipy.ndimage.grey_erosion(BrainMaskWeightedIMGFallback, footprint = numpy.ones((3, 3, 3)))
BrainMaskWeightedIMGFallback = scipy.ndimage.gaussian_filter(BrainMaskWeightedIMGFallback, sigma = [8 / x for x in BrainMaskNII.header.get_zooms()])
BrainMaskWeightedIMGFallback = scipy.ndimage.grey_dilation(BrainMaskWeightedIMGFallback, footprint = numpy.ones((3, 3, 3)))
BrainMaskWeightedIMGFallback = scipy.ndimage.grey_dilation(BrainMaskWeightedIMGFallback, footprint = numpy.ones((3, 3, 3)))
BrainMaskWeightedIMGFallback = scipy.ndimage.grey_dilation(BrainMaskWeightedIMGFallback, footprint = numpy.ones((3, 3, 3)))
BrainMaskWeightedIMGFallback = scipy.ndimage.grey_dilation(BrainMaskWeightedIMGFallback, footprint = numpy.ones((3, 3, 3)))
BrainMaskWeightedIMGFallback = scipy.ndimage.grey_dilation(BrainMaskWeightedIMGFallback, footprint = numpy.ones((3, 3, 3)))
BrainMaskWeightedIMGFallback = scipy.ndimage.grey_dilation(BrainMaskWeightedIMGFallback, footprint = numpy.ones((3, 3, 3)))
BrainMaskWeightedIMGFallback = scipy.ndimage.grey_dilation(BrainMaskWeightedIMGFallback, footprint = numpy.ones((3, 3, 3)))
BrainMaskWeightedIMGFallback = scipy.ndimage.grey_dilation(BrainMaskWeightedIMGFallback, footprint = numpy.ones((3, 3, 3)))
BrainMaskWeightedIMGFallback = scipy.ndimage.grey_dilation(BrainMaskWeightedIMGFallback, footprint = numpy.ones((3, 3, 3)))
BrainMaskWeightedIMGFallback = scipy.ndimage.grey_dilation(BrainMaskWeightedIMGFallback, footprint = numpy.ones((3, 3, 3)))
BrainMaskWeightedIMGFallback = scipy.ndimage.grey_dilation(BrainMaskWeightedIMGFallback, footprint = numpy.ones((3, 3, 3)))
BrainMaskWeightedIMGFallback = scipy.ndimage.grey_dilation(BrainMaskWeightedIMGFallback, footprint = numpy.ones((3, 3, 3)))
NewNII = nibabel.Nifti1Image(BrainMaskWeightedIMGFallback, BrainMaskNII.affine, header = BrainMaskNII.header)
nibabel.save(NewNII, ReferenceVolume + "_brain_mask_weighted_fallback.nii.gz")
NewNII = nibabel.Nifti1Image(BrainMaskWeightedIMG, BrainMaskNII.affine, header = BrainMaskNII.header)
nibabel.save(NewNII, ReferenceVolume + "_brain_mask_weighted.nii.gz")
#:wprint(ReferenceVolume)
#quit()
#shutil.copy(ReferenceVolume + "_brain_mask_weighted_fallback.nii.gz", "brain_mask_weighted_fallback.nii.gz")
#shutil.copy(ReferenceVolume + "_brain_mask_weighted.nii.gz", "brain_mask_weighted.nii.gz")

outputFiles = list()

if parallelMode == True:
    parallelSwitch = '-j+0'
    parallelPipe = subprocess.Popen(['parallel', parallelSwitch], stdin = subprocess.PIPE, stdout = subprocess.PIPE)

for curVolume in range(NumTotalVolumes):
    curMovingFile = os.path.join(TempDir, "tmp" + str(curVolume).zfill(4))
    TransformFileName = curMovingFile + '_to_tmp0000.mat'
    curOutputFile = curMovingFile + "_to_tmp0000"
    #CMD = ['bet', curMovingFile, curMovingFile + "_brain", '-f', '0.2', '-m', ';', 'flirt', '-in', curMovingFile, '-inweight', curMovingFile + "_brain_mask", '-ref', ReferenceVolume, '-omat', TransformFileName, '-out', curMovingFile + "_to_tmp0000", '-nosearch', '-bins', '256', '-paddingsize', '1', '-refweight', ReferenceVolume + "_brain_mask_weighted", '-cost', 'mutualinfo']
    CMD = ['flirt', '-in', curMovingFile, '-ref', ReferenceVolume, '-omat', TransformFileName, '-out', curMovingFile + "_to_tmp0000", '-nosearch', '-bins', '256', '-paddingsize', '1', '-refweight', ReferenceVolume + "_brain_mask_weighted", '-cost', 'mutualinfo']
    #print(" ".join(CMD))
    outputFiles.append(curOutputFile)

    #subprocess.call(['flirt', '-in', curMovingFile, '-ref', ReferenceVolumePlain, '-out', curMovingFile + "to_tmp0000", '-omat', TransformFileName, '-nosearch', '-paddingsize', '1'])
    #print " ".join(CMD) + "\n"
    if parallelMode == True:
        parallelPipe.stdin.write((" ".join(CMD) + "\n").encode('utf-8'))
    else:
        subprocess.call(CMD)

if parallelMode == True:
    parallelPipe.communicate()
    parallelPipe.stdin.close()

Iteration = 1;

InitialScales = None
InitialTranslation = None

while True:
    WeightedFallback = list()

    for curVolume in range(NumTotalVolumes):
        curMovingFile = os.path.join(TempDir, "tmp" + str(curVolume).zfill(4))
        TransformFileName = curMovingFile + '_to_tmp0000.mat'
        TransformFileFID = open(TransformFileName, 'r')
        #dialect = csv.Sniffer().sniff(TransformFileFID.read(1024))
        #TransformFileFID.seek(0)
        #TransformFileReader = csv.reader(TransformFileFID, dialect = dialect)
        Transform = numpy.zeros((4, 4))    
        
        curRow = 0
        for curLine in TransformFileFID:
            #Transforms[curRow, :, curVolume] = [row[0], row[1], row[2], row[3]]
            #print numpy.array([row[0], row[1], row[2], row[3]])
            row = curLine.split()
            for i in range(len(row)):
                row[i] = float(row[i])
            Transform[curRow, :] = [row[0], row[1], row[2], row[3]]
            curRow = curRow + 1
        TransformFileFID.close()
        (rotmat, skew, scales, transl, angles) = FLIRT.fsl_decomp_aff(numpy.matrix(Transform))
        if InitialTranslation is None:
            InitialTranslation = numpy.array(transl)
            TranslationDistance = 0
        else:
            TranslationDistance = numpy.sqrt(numpy.sum(numpy.multiply(InitialTranslation - transl, InitialTranslation - transl)))
            #print str(curVolume) + " " + str(TranslationDistance) + " " + str(angles.T) + " " + str(scales.diagonal()) 
        curScales = scales.diagonal() 
        # focus on the x and y scaling
        if numpy.any(numpy.logical_or(curScales[0:2] < 0.95, curScales[0:2] > 1.05)) or TranslationDistance > 10:
            WeightedFallback.append(curVolume)
    
    if len(WeightedFallback) > 0 and Iteration < 10:
    #if True == False:
        #print "Initial mask failed, iteration " + str(Iteration) + ", trying fallback on: " + " ".join([str(i) for i in WeightedFallback])
        if parallelMode == True:
            parallelPipe = subprocess.Popen(['parallel', parallelSwitch], stdin = subprocess.PIPE, stdout = subprocess.PIPE)

        for curVolume in WeightedFallback:
            curMovingFile = os.path.join(TempDir, "tmp" + str(curVolume).zfill(4))
            TransformFileName = curMovingFile + '_to_tmp0000.mat'
            curOutputFile = curMovingFile + "_to_tmp0000"
            CMD = ['flirt', '-in', curMovingFile, '-ref', ReferenceVolume, '-omat', TransformFileName, '-out', curMovingFile + "_to_tmp0000", '-nosearch', '-paddingsize', '1', '-refweight', ReferenceVolume + "_brain_mask_weighted_fallback", '-cost', 'mutualinfo']

            #subprocess.call(['flirt', '-in', curMovingFile, '-ref', ReferenceVolumePlain, '-out', curMovingFile + "to_tmp0000", '-omat', TransformFileName, '-nosearch', '-paddingsize', '1'])
            #print " ".join(CMD) + "\n"
            if parallelMode == True:
                parallelPipe.stdin.write((" ".join(CMD) + "\n").encode('utf-8'))
            else:
                subprocess.call(CMD)
            
        if parallelMode == True:
            parallelPipe.communicate()
            parallelPipe.stdin.close()
        # dilate the fallback in preparation for the next iteration
        subprocess.call(['fslmaths', ReferenceVolume + "_brain_mask_weighted_fallback", '-dilF', ReferenceVolume + "_brain_mask_weighted_fallback", '-odt', 'float'])
        Iteration = Iteration + 1
    else:
        os.unlink(ReferenceVolume + "_brain_mask_weighted_fallback.nii.gz")
        break

CMD = ['fslmerge', '-a', OutputFile]

CMD.extend(outputFiles)

subprocess.call(CMD)

if not OutputMATFile is None:
    fp = open(OutputMATFile, 'w')
    for i in outputFiles:
        (head, tail) = os.path.split(i)
        fp.write(tail)
        fp.write('\n')
        CurFP = open(i + '.mat', 'r')
        for curLine in CurFP:
            fp.write(curLine)
        CurFP.close()
    fp.close()

shutil.rmtree(TempDir)
