cmtk

@requires_cmtk
def cmtk_version(as_string=False):
    """Get CMTK version."""
    p = subprocess.run([_cmtkbin / 'streamxform', '--version'],
                       capture_output=True)
    version = p.stdout.decode('utf-8').rstrip()

    if as_string:
        return version
    else:
        return tuple(int(v) for v in version.split('.'))

def find_cmtkbin(tool: str = 'streamxform') -> str:
    """Find directory with CMTK binaries."""
    for path in _search_path:
        path = pathlib.Path(path)
        if not path.is_dir():
            continue

        try:
            return next(path.glob(tool)).resolve().parent            
        except StopIteration:
            continue
        except BaseException:
            raise

def parse_target_specs(target):
    """Parse target specs into argument that can be passed to CMTK."""
    # Note to self: this function should also deal with VoxelNeurons and NRRD filepaths
    # For NRRD filepaths: we need to add an empty "--" before the filepath (I think)

    from .templates import TemplateBrain

    assert isinstance(target, (str, TemplateBrain, np.ndarray, list, tuple))

    if isinstance(target, str):
        from . import registry 
        target = registry.find_template(target)

    if isinstance(target, TemplateBrain):
        specs = list(target.dims) + list(target.voxdims)
        # Note to self: need to check TemplateBrain (and flybrains) consistent definition of 
        # dims, voxdims and origin (maybe even add origin)

    # At this point we expect specs to be an iterable      
    specs = np.asarray(target)
    assert len(specs) in (6, 9), f"Target specs must be of length 6 or 9, got {len(specs)}"    
    target = "--target-grid "
    target += ",".join(map(str, specs[:3].astype(int)))  # Number of voxels (must be integer)
    target += ":"
    target += ",".join(map(str, specs[3:].astype(float))) # Voxel size (can be float)
    if len(specs) == 9:
        target += ":"
        target += ",".join(map(str, specs[6:].astype(float))) # Origin (can be float)

    return target

def requires_cmtk(func):
    """Check if CMTK is available."""
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        if not _cmtkbin:
            raise ValueError("Cannot find CMTK. Please install from "
                             "http://www.nitrc.org/projects/cmtk and "
                             "make sure that it is your path!")
        return func(*args, **kwargs)
    return wrapper

@requires_cmtk
def xform_cmtk(points: np.ndarray, transforms, inverse: bool = False,
               affine_fallback: bool = False, **kwargs) -> np.ndarray:
    """Xform 3d coordinates.

    Parameters
    ----------
    points :            (N, 3) array | pandas.DataFrame
                        Points to transform. DataFrame must have x/y/z columns.
    transforms :        filepath(s) | CMTKtransform(s)
                        Either filepath to CMTK transform or `CMTKtransform`.
                        Multiple regs must be given as list and will be applied
                        sequentially in the order provided.
    inverse :           bool | list thereof
                        Whether to invert transforms. If single boolean will
                        apply to all transforms. Can also provide `inverse` as
                        list of booleans.
    affine_fallback :   bool
                        If True, points that failed to transform during warping
                        transform will be transformed using only the affine
                        transform.

    Returns
    -------
    pointsxf :          (N, 3) numpy.ndarray
                        Transformed points. Will contain `np.nan` for points
                        that did not transform.

    """
    transforms = list(utils.make_iterable(transforms))

    if isinstance(inverse, bool):
        inverse = [inverse] * len(transforms)

    directions = ['forward' if not i else 'inverse' for i in inverse]

    for i, r in enumerate(transforms):
        if not isinstance(r, CMTKtransform):
            if not isinstance(r, (str, pathlib.Path)):
                raise TypeError('`reg` must be filepath or CMTKtransform')
            transforms[i] = CMTKtransform(r, directions=directions[i])

    # Combine all transforms into a sequence of transforms
    seq = TransformSequence(*transforms)

    # Transform points
    xf = seq.xform(points)

    # If requested, try again with affine only for points that failed to xform
    if affine_fallback:
        isnan = np.any(np.isnan(xf), axis=1)
        if np.any(isnan):
            xf[isnan] = seq.xform(points[isnan], affine_only=True)

    return xf