import numpy
import argparse

parser = argparse.ArgumentParser(
  description='Shift particle centers in a TRF file',
  epilog='Written by Dustin R. Morado 28.09.2016')
parser.add_argument('input', help='Input trf file', metavar='IN.trf')
parser.add_argument('output', help='Output trf file', metavar='OUT.trf')

parser.add_argument(
    '--shiftX', '-x',
    type=float,
    help='Shift in X',
    default=0.0)

parser.add_argument(
    '--shiftY', '-y',
    type=float,
    help='Shift in Y',
    default=0.0)

parser.add_argument(
    '--shiftZ', '-z',
    type=float,
    help='Shift in Z',
    default=0.0)

args = parser.parse_args()

# We take in the shift relative to the reference structure
referenceShift = numpy.array(
    [[args.shiftX],
     [args.shiftY],
     [args.shiftZ]],
    numpy.float_)

with open(args.input) as inputTRF, open(args.output, 'w') as outputTRF:
  for initialTransformLine in inputTRF:
    initialTransform = initialTransformLine.split()
    # Initialize new transform line by copying subset
    finalTransformLine = initialTransform[0]

    initialCoordinates = numpy.array(
        [[initialTransform[1]],
         [initialTransform[2]],
         [initialTransform[3]]],
        numpy.int_)

    initialDisplacement = numpy.array(
        [[initialTransform[4]],
         [initialTransform[5]],
         [initialTransform[6]]],
        numpy.float_)

    # The transform matrix describes the rotation of the particle coordinate
    # system to the reference system
    initialRotationMatrix = numpy.array(
        [[initialTransform[7], initialTransform[8], initialTransform[9]],
         [initialTransform[10], initialTransform[11], initialTransform[12]],
         [initialTransform[13], initialTransform[14], initialTransform[15]]],
        numpy.float_)

    # We are interested in the inverse transform above, which for rotation
    # matrices is just the transform
    inverseRotationMatrix = initialRotationMatrix.T

    initialParticleCenter = initialCoordinates + initialDisplacement

    # We apply the inverse rotation transform to our shift vector to describe
    # the requested shift in the particle's coordinate system
    particleShift = inverseRotationMatrix.dot(referenceShift)

    finalParticleCenter = initialParticleCenter + particleShift
    finalDisplacement, finalCoordinates = numpy.modf(finalParticleCenter)
    finalCoordinates = finalCoordinates.astype(numpy.int_)

    finalTransformLine += " {:d} {:d} {:d} ".format(
        finalCoordinates[0,0], finalCoordinates[1,0], finalCoordinates[2,0])
    finalTransformLine += "{: 14.10f} {: 14.10f} {: 14.10f} ".format(
        finalDisplacement[0,0], finalDisplacement[1,0], finalDisplacement[2,0])
    finalTransformLine += "{:s} {:s} {:s} ".format(
        initialTransform[7], initialTransform[8], initialTransform[9])
    finalTransformLine += "{:s} {:s} {:s} ".format(
        initialTransform[10], initialTransform[11], initialTransform[12])
    finalTransformLine += "{:s} {:s} {:s}\n".format(
        initialTransform[13], initialTransform[14], initialTransform[15])
    outputTRF.write(finalTransformLine)