transforms

class AffineTransform(BaseTransform):
    """Affine transformation of 3D spatial data.

    Parameters
    ----------
    matrix :        (4, 4) np.ndarray
                    Affine matrix.

    Examples
    --------
    A simple scaling transform

    >>> from navis import transforms
    >>> import numpy as np
    >>> M = np.diag([1e3, 1e3, 1e3, 1])
    >>> tr = transforms.affine.AffineTransform(M)
    >>> points = np.array([[0, 0, 0], [1, 1, 1]])
    >>> tr.xform(points)
    array([[   0.,    0.,    0.],
           [1000., 1000., 1000.]])

    """

    def __init__(self, matrix: np.ndarray, direction: str = 'forward'):
        """Initialize transform."""
        assert direction in ('forward', 'inverse')

        self.matrix = matrix

        if direction == 'inverse':
            self.matrix = np.linalg.inv(self.matrix)

    def __eq__(self, other: 'AffineTransform') -> bool:
        """Implements equality comparison."""
        if isinstance(other, AffineTransform):
            if np.all(self.matrix == other.matrix):
                return True
        return False

    def __neg__(self):
        """Invert direction."""
        x = self.copy()

        # Invert affine matrix
        x.matrix = np.linalg.inv(x.matrix)

        return x

    def copy(self) -> 'AffineTransform':
        """Return copy of transform."""
        # Attributes not to copy
        no_copy = []
        # Generate new empty transform
        x = self.__class__(None)
        # Override with this neuron's data
        x.__dict__.update({k: copy.copy(v) for k, v in self.__dict__.items() if k not in no_copy})

        return x

    def xform(self, points: np.ndarray, invert: bool = False) -> np.ndarray:
        """Apply transform to points.

        Parameters
        ----------
        points :        np.ndarray
                        (N, 3) array of x/y/z locations.
        invert :        bool
                        If True, will invert the transform.

        Returns
        -------
        pointsxf :      np.ndarray
                        The transformed points.

        """
        points = np.asarray(points)

        if points.ndim != 2 and points.shape[1] != 3:
            raise ValueError('`points` must be of shape (N, 3)')

        # Add a fourth column to points
        points_mat = np.ones((points.shape[0], 4))
        points_mat[:, :3] = points

        # Apply transform
        if not invert:
            mat = self.matrix
        else:
            mat = np.linalg.inv(self.matrix)

        return np.dot(mat, points_mat.T).T[:, :3]

class AliasTransform(BaseTransform):
    """Helper transform that simply passes points through.

    Useful for defining aliases.
    """

    def __init__(self):
        """Initialize."""
        pass

    def __neg__(self) -> 'AliasTransform':
        """Invert transform."""
        return self.copy()

    def __eq__(self, other):
        """Check if the same."""
        if isinstance(other, AliasTransform):
            True
        return False

    def copy(self):
        """Return copy."""
        x = AliasTransform()
        x.__dict__.update(self.__dict__)
        return x

    def xform(self, points: np.ndarray) -> np.ndarray:
        """Pass through.

        Note that the returned points are NOT a copy but the originals.
        """
        return points

class CMTKtransform(BaseTransform):
    """CMTK transforms of 3D spatial data.

    Requires [CMTK](https://www.nitrc.org/projects/cmtk/) to be installed.

    Parameters
    ----------
    regs :          str | list of str
                    Path(s) to CMTK transformations(s).
    directions :    "forward" | "inverse" | list thereof
                    Direction of transformation. Must provide one direction per
                    `reg`.
    threads :       int, optional
                    Number of threads to use.

    Examples
    --------
    >>> from navis import transforms
    >>> tr = transforms.cmtk.CMTKtransform('/path/to/CMTK_directory.list')
    >>> tr.xform(points) # doctest: +SKIP

    """

    def __init__(self, regs: list, directions: str = 'forward', threads: int = None):
        self.directions = list(utils.make_iterable(directions))
        for d in self.directions:
            assert d in ('forward', 'inverse'), ('`direction` must be "foward"'
                                                 f'or "inverse", not "{d}"')

        self.regs = list(utils.make_iterable(regs))
        self.command = 'streamxform'
        self.threads = threads

        if len(directions) == 1 and len(regs) >= 1:
            directions = directions * len(regs)

        if len(self.regs) != len(self.directions):
            raise ValueError('Must provide one direction per regs')

    def __eq__(self, other: 'CMTKtransform') -> bool:
        """Implement equality comparison."""
        if isinstance(other, CMTKtransform):
            if len(self) == len(other):
                if all([self.regs[i] == other.regs[i] for i in range(len(self))]):
                    if all([self.directions[i] == other.directions[i] for i in range(len(self))]):
                        return True
        return False

    def __len__(self) -> int:
        return len(self.regs)

    def __neg__(self) -> 'CMTKtransform':
        """Invert direction."""
        x = self.copy()

        # Swap directions
        x.directions = [{'forward': 'inverse',
                         'inverse': 'forward'}[d] for d in x.directions]

        # Reverse order
        x.regs = x.regs[::-1]
        x.directions = x.directions[::-1]

        return x

    def __str__(self):
        return self.__repr__()

    def __repr__(self):
        return f'CMTKtransform with {len(self)} transform(s)'

    @staticmethod
    def from_file(filepath: str, **kwargs) -> 'CMTKtransform':
        """Generate CMTKtransform from file.

        Parameters
        ----------
        filepath :  str
                    Path to CMTK transform.
        **kwargs
                    Keyword arguments passed to CMTKtransform.__init__

        Returns
        -------
        CMTKtransform

        """
        defaults = {'directions': 'forward'}
        defaults.update(kwargs)
        return CMTKtransform(str(filepath), **defaults)

    def make_args(self, affine_only: bool = False) -> list:
        """Generate arguments passed to subprocess."""
        # Generate the arguments
        # The actual command (i.e. streamxform)
        args = [str(_cmtkbin / self.command)]

        if affine_only:
            args.append('--affine-only')

        if self.threads:
            args.append(f'--threads {int(self.threads)}')

        # Add the regargs
        args += self.regargs

        return args

    @property
    def regargs(self) -> list:
        """Generate regargs."""
        regargs = []
        for i, (reg, dir) in enumerate(zip(self.regs, self.directions)):
            if dir == 'inverse':
                # For the first transform we need to prefix "--inverse" with
                # a solitary "--"
                if i == 0:
                    regargs.append('--')
                regargs.append('--inverse')
            # Note no double quotes!
            regargs.append(f'{reg}')
        return regargs

    def append(self, transform: 'CMTKtransform', direction: str = None):
        """Add another transform.

        Parameters
        ----------
        transform :     str | CMTKtransform
                        Either another CMTKtransform or filepath to registration.
        direction :     "forward" | "inverse"
                        Only relevant if transform is filepath.

        """
        if isinstance(transform, CMTKtransform):
            if self.command != transform.command:
                raise ValueError('Unable to merge CMTKtransforms using '
                                 'different commands.')

            self.regs += transform.regs
            self.directions += transform.directions
        elif isinstance(transform, str):
            if not direction:
                raise ValueError('Must provide direction along with new transform')
            self.regs.append(transform)
            self.directions.append(direction)
        else:
            raise NotImplementedError(f'Unable to append {type(transform)} to {type(self)}')

    def check_if_possible(self, on_error: str = 'raise'):
        """Check if this transform is possible."""
        if not _cmtkbin:
            msg = 'Folder with CMTK binaries not found. Make sure the ' \
                  'directory is in your PATH environment variable.'
            if on_error == 'raise':
                raise BaseException(msg)
            return msg
        for r in self.regs:
            if not os.path.isdir(r) and not os.path.isfile(r):
                msg = f'Registration {r} not found.'
                if on_error == 'raise':
                    raise BaseException(msg)
                return msg

    def copy(self) -> 'CMTKtransform':
        """Return copy."""
        # Attributes not to copy
        no_copy = []
        # Generate new empty transform
        x = self.__class__(None)
        # Override with this neuron's data
        x.__dict__.update({k: copy.copy(v) for k, v in self.__dict__.items() if k not in no_copy})

        return x

    def parse_cmtk_output(self, output: str, fail_value=np.nan) -> np.ndarray:
        r"""Parse CMTK output.

        Briefly, CMTK output will be a byte literal like this:

            b'311 63 23 \n275 54 25 \n'

        In case of failed transforms we will get something like this where the
        original coordinates are returned with a "FAILED" flag

            b'343 72 23 \n-10 -10 -10 FAILED \n'

        Parameter
        ---------
        output :        tuple of (b'', None)
                        Stdout of CMTK call.
        fail_value
                        Value to use for points that failed to transform. By
                        default we use `np.nan`.

        Returns
        -------
        pointsxf :      (N, 3) numpy array
                        The parse transformed points.

        """
        # The original stout is tuple where we care only about the second one
        if isinstance(output, tuple):
            output = output[0]

        pointsx = []
        # Split string into rows - lazily using a generator
        for row in (x.group(1) for x in re.finditer(r"(.*?) \n",
                                                    output.decode())):
            # Split into values
            values = row.split(' ')

            # If this point failed
            if len(values) != 3:
                values = [fail_value] * 3
            else:
                values = [float(v) for v in values]

            pointsx.append(values)

        return np.asarray(pointsx)

    def xform(self, points: np.ndarray,
              affine_only: bool = False,
              affine_fallback: bool = False) -> np.ndarray:
        """Xform data.

        Parameters
        ----------
        points :            (N, 3) numpy array | pandas.DataFrame
                            Points to xform. DataFrame must have x/y/z columns.
        affine_only :       bool
                            Whether to apply only the non-rigid affine
                            transform. This is useful if points are outside
                            the deformation field and would therefore not
                            transform properly.
        affine_fallback :   bool
                            If True and some points did not transform during the
                            non-rigid part of the transformation, we will apply
                            only the affine transformation to those points.

        Returns
        -------
        pointsxf :      (N, 3) numpy array
                        Transformed points. Points that failed to transform will
                        be `np.nan`.

        """
        self.check_if_possible(on_error='raise')

        if isinstance(points, pd.DataFrame):
            # Make sure x/y/z columns are present
            if np.any([c not in points for c in ['x', 'y', 'z']]):
                raise ValueError('points DataFrame must have x/y/z columns.')
        elif isinstance(points, np.ndarray) and points.ndim == 2 and points.shape[1] == 3:
            points = pd.DataFrame(points, columns=['x', 'y', 'z'])
        else:
            raise TypeError('`points` must be numpy array of shape (N, 3) or '
                            'pandas DataFrame with x/y/z columns')

        # Generate the result
        args = self.make_args(affine_only=affine_only)
        proc = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE)

        # Pipe in the points
        points_str = points[['x', 'y', 'z']].to_string(index=False,
                                                       header=False)

        # Do not use proc.stdin.write to avoid output buffer becoming full
        # before we finish piping in stdin.
        #proc.stdin.write(points_str.encode())
        #output = proc.communicate()

        # Read out results
        # This is equivalent to e.g.:
        # $ streamxform -args <<< "10, 10, 10"
        output = proc.communicate(input=points_str.encode())

        # If no output, something went wrong
        if not output[0]:
            raise utils.CMTKError('CMTK produced no output. Check points?')

        # Xformed points
        xf = self.parse_cmtk_output(output, fail_value=np.nan)

        # Check if any points not xformed
        if affine_fallback and not affine_only:
            not_xf = np.any(np.isnan(xf), axis=1)
            if np.any(not_xf):
                xf[not_xf] = self.xform(points.loc[not_xf], affine_only=True)

        return xf

    def xform_image(self,
                    im,
                    target,
                    out=None,
                    interpolation="linear",                    
                    verbose=False,
                    ):
        """Transform an image using CMTK's reformatx.

        Parameters 
        ----------
        im :        3D numpy array | filepath
                    The floating image to transform.
        target :    str | TemplateBrain | (Nx, Ny, Nz, dx, dy, dz) | (Nx, Ny, Nz, dx, dy, dz, Ox, Oy, Oz)
                    Defines the target image: dimensions in voxels (N), the voxel size (d) and optionally 
                    an origin (0) for the target image. Can be provided as a string (name of a template),
                    a TemplateBrain object, a tuple/list/array with the target specs.
        out :       str, optional
                    The filepath to save the transformed image. If None (default), will return the 
                    transformed image as np.ndarray.
        interpolation : "linear" | "nn" | "cubic" | "pv" | "sinc-cosine" | "sinc-hamming"
                    The interpolation method to use.
        verbose :   bool
                    Whether to print CMTK output.

        Returns
        -------
        np.ndarray | None
                    If out is None, returns the transformed image as np.ndarray. Otherwise, None.

        """
        assert interpolation in ("linear", "nn", "cubic", "pv", "sinc-cosine", "sinc-hamming")

        # `reformatx` expects this format:
        # ./reformatx --floating {INPUT_FILE} -o {OUTPUT_FILE} {REFERENCE_SPECS} {TRANSFORMS} 
        # where:
        # - {INPUT_FILE} is the image to transform
        # - {OUTPUT_FILE} is where the output will be saved
        # - {REFERENCE_SPECS} defines the target space; this needs to be eitheran NRRD
        #   file from which CMTK can extract the target grid or the actual specs: 
        #   "--target-grid Nx,Ny,Nz:dX,dY,dZ:[Ox,Oy,Oz]" where N is the number of
        #   voxels in each dimension and d is the voxel size. The optional O is the
        #   origin of the image. If not provided, it is assumed to be (0, 0, 0).
        # - {TRANSFORMS} are the CMTK transform(s) to apply; prefix with "--inverse" to invert
        # Below command works to convert JFRC2 to FCWB:
        # /opt/local/lib/cmtk/bin/reformatx --verbose --floating JFRC2.nrrd -o JFRC2_xf.nrrd FCWB.nrrd --inverse /Users/philipps/flybrain-data/BridgingRegistrations/JFRC2_FCWB.list
        # This took XX minutes - should check if that is actually faster than the look-up approach we 
        # use in `images.py`

        target_specs = parse_target_specs(target)

        to_remove = []
        if isinstance(im, (str, pathlib.Path)):
            floating = pathlib.Path(im)
            if not im.is_file():
                raise ValueError(f"Image file not found: {im}")
        elif isinstance(im, np.ndarray):
            assert im.ndim == 3
            # Save to temporary file
            with tempfile.NamedTemporaryFile(suffix=".nrrd", delete=False, delete_on_close=False) as tf:
                nrrd.write(tf.name, im)
                floating = tf.name
                to_remove.append(tf.name)
        else:
            raise ValueError(f"Invalid image type: {type(im)}")

        if out is None:
            outfile = tempfile.NamedTemporaryFile(suffix=".nrrd", delete=False, delete_on_close=False).name
            to_remove.append(outfile)
        elif isinstance(out, (str, pathlib.Path)):
            outfile = pathlib.Path(out).resolve()
        else:
            raise ValueError(f"Invalid output type: {type(out)}")

        # Compile the command 
        args = [str(_cmtkbin / 'reformatx')]
        args += [f'-o {outfile}']
        args += [f'--floating {floating}']
        args += [f'--{interpolation}']
        args += [target_specs]

        # Add the regargs
        args += self.regargs      

        try:
            # run the binary
            # avoid resourcewarnings with null
            with open(os.devnull, "w") as devnull:
                startupinfo = None
                if platform.system() == "Windows":
                    startupinfo = subprocess.STARTUPINFO()
                    startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
                if verbose:
                    # in debug mode print the output
                    stdout = None
                else:
                    stdout = devnull

                if verbose:
                    config.logger.info("executing: {}".format(" ".join(args)))
                check_call(
                    args,
                    stdout=stdout,
                    stderr=subprocess.STDOUT,
                    startupinfo=startupinfo,
                )  

            if out is None:
                # Return transformed image
                return nrrd.read(outfile)
            elif verbose:
                config.logger.info(f"Transformed image saved to {outfile}")
        except BaseException:
            raise 
        finally:
            # Clean up temporary files
            for f in to_remove:
                os.remove(f) 

class ElastixTransform(BaseTransform):
    """Elastix transforms of 3D spatial data.

    Requires [Elastix](https://github.com/SuperElastix/elastix/). Based on
    code by Jasper Phelps (<https://github.com/jasper-tms/pytransformix>).

    Note that elastix transforms can not be inverted!

    Parameters
    ----------
    file :              str
                        Filepath to elastix transformation file.
    copy_files :        filepath | list, optional
                        Any files that need to be copied into the temporary
                        directory where we perform the transform. These are
                        typically files supplemental to the main transform
                        file (e.g. defining an additional affine transform).

    Examples
    --------
    >>> from navis import transforms
    >>> tr = transforms.ElastixTransform('/path/to/transform/transform')
    >>> tr.xform(points) # doctest: +SKIP

    """

    def __init__(self, file: str, copy_files=[]):
        self.file = pathlib.Path(file)
        self.copy_files = copy_files

    def __eq__(self, other: "ElastixTransform") -> bool:
        """Implement equality comparison."""
        if isinstance(other, ElastixTransform):
            if self.file == other.file:
                return True
        return False

    def check_if_possible(self, on_error: str = "raise"):
        """Check if this transform is possible."""
        if not _elastixbin:
            msg = (
                "Folder with elastix binaries not found. Make sure the "
                "directory is in your PATH environment variable."
            )
            if on_error == "raise":
                raise BaseException(msg)
            return msg
        if not self.file.is_file():
            msg = f"Transformation file {self.file} not found."
            if on_error == "raise":
                raise BaseException(msg)
                return msg

    def copy(self) -> "ElastixTransform":
        """Return copy."""
        # Attributes not to copy
        no_copy = []
        # Generate new empty transform
        x = self.__class__(self.file)
        # Override with this neuron's data
        x.__dict__.update(
            {k: copy.copy(v) for k, v in self.__dict__.items() if k not in no_copy}
        )

        return x

    def write_input_file(self, points, filepath):
        """Write a numpy array in format required by transformix."""
        with open(filepath, "w") as f:
            f.write("point\n{}\n".format(len(points)))
            for x, y, z in points:
                f.write(f"{x:f} {y:f} {z:f}\n")

    def read_output_file(self, filepath) -> np.ndarray:
        """Load output file.

        Parameter
        ---------
        filepath :      str
                        Filepath to output file.

        Returns
        -------
        pointsxf :      (N, 3) numpy array
                        The parse transformed points.

        """
        points = []
        with open(filepath, "r") as f:
            for line in f.readlines():
                output = line.split("OutputPoint = [ ")[1].split(" ]")[0]
                points.append([float(i) for i in output.split(" ")])
        return np.array(points)

    def xform(self, points: np.ndarray, return_logs=False) -> np.ndarray:
        """Xform data.

        Parameters
        ----------
        points :        (N, 3) numpy array | pandas.DataFrame
                        Points to xform. DataFrame must have x/y/z columns.
        return_logs :   bool
                        If True, will return logs instead of transformed points.
                        Really only useful for debugging.

        Returns
        -------
        pointsxf :      (N, 3) numpy array
                        Transformed points.

        """
        self.check_if_possible(on_error="raise")

        if isinstance(points, pd.DataFrame):
            # Make sure x/y/z columns are present
            if np.any([c not in points for c in ["x", "y", "z"]]):
                raise ValueError("points DataFrame must have x/y/z columns.")
            points = points[["x", "y", "z"]].values
        elif not (
            isinstance(points, np.ndarray) and points.ndim == 2 and points.shape[1] == 3
        ):
            raise TypeError(
                "`points` must be numpy array of shape (N, 3) or "
                "pandas DataFrame with x/y/z columns"
            )

        # Everything happens in a temporary directory
        with tempfile.TemporaryDirectory() as tempdir:
            p = pathlib.Path(tempdir)

            # If required, copy additional files into the temporary directory
            if self.copy_files:
                for f in make_iterable(self.copy_files):
                    _ = pathlib.Path(shutil.copy(f, p))

            # Write points to file
            in_file = p / "inputpoints.txt"
            self.write_input_file(points, in_file)

            out_file = p / "outputpoints.txt"

            # Prepare the command
            command = [
                _elastixbin / "transformix",
                "-out",
                str(p),
                "-tp",
                str(self.file),
                "-def",
                str(in_file),
            ]

            # Keep track of current working directory
            cwd = os.getcwd()
            try:
                # Change working directory to the temporary directory
                # This is apparently required because elastix stupidly expects
                # any secondary transform files to be in the current directory
                # (as opposed to where the main transform is)
                os.chdir(p)
                # Run the actual transform
                proc = subprocess.run(command, stdout=subprocess.PIPE)
            except BaseException:
                raise
            finally:
                # This makes sure we land on our feet even in case of an error
                os.chdir(cwd)

            if return_logs:
                logfile = p / "transformix.log"
                if not logfile.is_file():
                    raise FileNotFoundError("No log file found.")
                with open(logfile) as f:
                    logs = f.read()
                return logs

            if not out_file.is_file():
                raise FileNotFoundError(
                    "Elastix transform did not produce any "
                    f"output:\n {proc.stdout.decode()}"
                )

            points_xf = self.read_output_file(out_file)

        return points_xf

class FunctionTransform(BaseTransform):
    """Apply custom function as transform.

    Parameters
    ----------
    func :      callable
                Function that accepts and returns an (N, 3) array.

    """

    def __init__(self, func):
        """Initialize."""
        if not callable(func):
            raise TypeError('`func` must be callable')
        self.func = func

    def __eq__(self, other):
        """Check if the same."""
        if not isinstance(other, FunctionTransform):
            return False
        if self.func != other.func:
            return False
        return True

    def copy(self):
        """Return copy."""
        x = self.__class__(self.func)
        x.__dict__.update(self.__dict__)
        return x

    def xform(self, points: np.ndarray) -> np.ndarray:
        """Xform data.

        Parameters
        ----------
        points :        (N, 3) numpy array | pandas.DataFrame
                        Points to xform. DataFrame must have x/y/z columns.

        Returns
        -------
        pointsxf :      (N, 3) numpy array
                        Transformed points.

        """
        if isinstance(points, pd.DataFrame):
            # Make sure x/y/z columns are present
            if np.any([c not in points for c in ['x', 'y', 'z']]):
                raise ValueError('points DataFrame must have x/y/z columns.')
            points = points[['x', 'y', 'z']].values
        elif not (isinstance(points, np.ndarray) and points.ndim == 2 and points.shape[1] == 3):
            raise TypeError('`points` must be numpy array of shape (N, 3) or '
                            'pandas DataFrame with x/y/z columns')
        return self.func(points.copy())