"""
============================================
Tractography Clustering - Available Features
============================================

This page lists available features that can be used by the tractography
clustering framework. For every feature a brief description is provided
explaining: what it does, when it's useful and how to use it. If you are not
familiar with the tractography clustering framework, read the
:ref:`clustering-framework` first.

.. contents:: Available Features
    :local:
    :depth: 1

**Note**:
All examples assume a function `get_streamlines` exists. We defined here a
simple function to do so. It imports the necessary modules and loads a small
streamline bundle.
"""


def get_streamlines():
    from dipy.data import get_fnames
    from dipy.io.streamline import load_tractogram
    from dipy.tracking.streamline import Streamlines

    fname = get_fnames('fornix')
    fornix = load_tractogram(fname, 'same',
                             bbox_valid_check=False).streamlines

    streamlines = Streamlines(fornix)
    return streamlines

###############################################################################
# .. _clustering-examples-IdentityFeature:
# 
# Identity Feature
# ================
# **What:** Instances of `IdentityFeature` simply return the streamlines
# unaltered.  In other words the features are the original data.
# 
# **When:** The QuickBundles algorithm requires streamlines to have the same
# number of points. If this is the case for your streamlines, you can tell
# QuickBundles to not perform resampling (see following example). The clustering
# should be faster than using the default behaviour of QuickBundles since it will
# require less computation (i.e. no resampling). However, it highly depends
# on the number of points streamlines have. By default, QuickBundles resamples
# streamlines so that they have 12 points each [Garyfallidis12]_.
# 
# *Unless stated otherwise, it is the default feature used by `Metric` objects
# in the clustering framework.*


from dipy.segment.clustering import QuickBundles
from dipy.segment.metric import AveragePointwiseEuclideanMetric
from dipy.segment.featurespeed import IdentityFeature

# Get some streamlines.
streamlines = get_streamlines()  # Previously defined.

# Make sure our streamlines have the same number of points.
from dipy.tracking.streamline import set_number_of_points
streamlines = set_number_of_points(streamlines, nb_points=12)

# Create an instance of `IdentityFeature` and tell metric to use it.
feature = IdentityFeature()
metric = AveragePointwiseEuclideanMetric(feature=feature)
qb = QuickBundles(threshold=10., metric=metric)
clusters = qb.cluster(streamlines)

print("Nb. clusters:", len(clusters))
print("Cluster sizes:", list(map(len, clusters)))

###############################################################################
# 
# ::
# 
#     Nb. clusters: 4
# 
#     Cluster sizes: [64, 191, 47, 1]
# 
# 
# .. _clustering-examples-ResampleFeature:
# 
# Resample Feature
# ================
# **What:** Instances of `ResampleFeature` resample streamlines to a
# predetermined number of points. The resampling is done on the fly such that
# there are no permanent modifications made to your streamlines.
# 
# **When:** The QuickBundles algorithm requires streamlines to have the same
# number of points. By default, QuickBundles uses `ResampleFeature` to resample
# streamlines so that they have 12 points each [Garyfallidis12]_. If you want to
# use a different number of points for the resampling, you should provide your
# own instance of `ResampleFeature` (see following example).
# 
# **Note:** Resampling streamlines has an impact on clustering results both in
# term of speed and quality. Setting the number of points too low will result in
# a loss of information about the shape of the streamlines. On the contrary,
# setting the number of points too high will slow down the clustering process.


from dipy.segment.clustering import QuickBundles
from dipy.segment.featurespeed import ResampleFeature
from dipy.segment.metric import AveragePointwiseEuclideanMetric

# Get some streamlines.
streamlines = get_streamlines()  # Previously defined.

# Streamlines will be resampled to 24 points on the fly.
feature = ResampleFeature(nb_points=24)
metric = AveragePointwiseEuclideanMetric(feature=feature)  # a.k.a. MDF
qb = QuickBundles(threshold=10., metric=metric)
clusters = qb.cluster(streamlines)

print("Nb. clusters:", len(clusters))
print("Cluster sizes:", list(map(len, clusters)))

###############################################################################
# 
# ::
# 
#     Nb. clusters: 4
# 
#     Cluster sizes: [64, 191, 44, 1]
# 
# .. _clustering-examples-CenterOfMassFeature:
# 
# Center of Mass Feature
# ======================
# **What:** Instances of `CenterOfMassFeature` compute the center of mass (also
# known as center of gravity) of a set of points. This is achieved by taking the
# mean of every coordinate independently (for more information see the
# `wiki page <https://en.wikipedia.org/wiki/Center_of_mass>`_).
# 
# **When:** This feature can be useful when you *only* need information about the
# spatial position of a streamline.
# 
# **Note:** The computed center is not guaranteed to be an existing point in the
# streamline.


import numpy as np
from dipy.viz import window, actor, colormap
from dipy.segment.clustering import QuickBundles
from dipy.segment.featurespeed import CenterOfMassFeature
from dipy.segment.metric import EuclideanMetric

# Enables/disables interactive visualization
interactive = False

# Get some streamlines.
streamlines = get_streamlines()  # Previously defined.


feature = CenterOfMassFeature()
metric = EuclideanMetric(feature)

qb = QuickBundles(threshold=5., metric=metric)
clusters = qb.cluster(streamlines)

# Extract feature of every streamline.
centers = np.asarray(list(map(feature.extract, streamlines)))

# Color each center of mass according to the cluster they belong to.
colormap = colormap.create_colormap(np.arange(len(clusters)))
colormap_full = np.ones((len(streamlines), 3))
for cluster, color in zip(clusters, colormap):
    colormap_full[cluster.indices] = color

# Visualization
scene = window.Scene()
scene.clear()
scene.SetBackground(0, 0, 0)
scene.add(actor.streamtube(streamlines, window.colors.white, opacity=0.05))
scene.add(actor.point(centers[:, 0, :], colormap_full, point_radius=0.2))
window.record(scene, n_frames=1, out_path='center_of_mass_feature.png',
              size=(600, 600))
if interactive:
    window.show(scene)

###############################################################################
# .. figure:: center_of_mass_feature.png
#    :align: center
# 
#    Showing the center of mass of each streamline colored according to
#    the QuickBundles results.
# 
# .. _clustering-examples-MidpointFeature:
# 
# Midpoint Feature
# ================
# **What:** Instances of `MidpointFeature` extract the middle point of a
# streamline. If there is an even number of points, the feature will then
# correspond to the point halfway between the two middle points.
# 
# **When:** This feature can be useful when you *only* need information about the
# spatial position of a streamline. This can also be an alternative to the
# `CenterOfMassFeature` if the point extracted must be on the streamline.


import numpy as np
from dipy.viz import window, actor, colormap
from dipy.segment.clustering import QuickBundles
from dipy.segment.featurespeed import MidpointFeature
from dipy.segment.metric import EuclideanMetric

# Get some streamlines.
streamlines = get_streamlines()  # Previously defined.

feature = MidpointFeature()
metric = EuclideanMetric(feature)

qb = QuickBundles(threshold=5., metric=metric)
clusters = qb.cluster(streamlines)

# Extract feature of every streamline.
midpoints = np.asarray(list(map(feature.extract, streamlines)))

# Color each midpoint according to the cluster they belong to.
colormap = colormap.create_colormap(np.arange(len(clusters)))
colormap_full = np.ones((len(streamlines), 3))
for cluster, color in zip(clusters, colormap):
    colormap_full[cluster.indices] = color

# Visualization
scene = window.Scene()
scene.clear()
scene.SetBackground(0, 0, 0)
scene.add(actor.point(midpoints[:, 0, :], colormap_full, point_radius=0.2))
scene.add(actor.streamtube(streamlines, window.colors.white, opacity=0.05))
window.record(scene, n_frames=1, out_path='midpoint_feature.png',
              size=(600, 600))
if interactive:
    window.show(scene)

###############################################################################
# .. figure:: midpoint_feature.png
#    :align: center
# 
#    Showing the middle point of each streamline colored according to the
#    QuickBundles results.
# 
# .. _clustering-examples-ArcLengthFeature:
# 
# ArcLength Feature
# =================
# **What:** Instances of `ArcLengthFeature` compute the length of a streamline.
# More specifically, this feature corresponds to the sum of the lengths of every
# streamline's segments.
# 
# **When:** This feature can be useful when you *only* need information about the
# length of a streamline.


import numpy as np
from dipy.viz import window, actor, colormap
from dipy.segment.clustering import QuickBundles
from dipy.segment.featurespeed import ArcLengthFeature
from dipy.segment.metric import EuclideanMetric

# Get some streamlines.
streamlines = get_streamlines()  # Previously defined.

feature = ArcLengthFeature()
metric = EuclideanMetric(feature)
qb = QuickBundles(threshold=2., metric=metric)
clusters = qb.cluster(streamlines)

# Color each streamline according to the cluster they belong to.
colormap = colormap.create_colormap(np.ravel(clusters.centroids))
colormap_full = np.ones((len(streamlines), 3))
for cluster, color in zip(clusters, colormap):
    colormap_full[cluster.indices] = color

# Visualization
scene = window.Scene()
scene.clear()
scene.SetBackground(0, 0, 0)
scene.add(actor.streamtube(streamlines, colormap_full))
window.record(scene, out_path='arclength_feature.png', size=(600, 600))
if interactive:
    window.show(scene)

###############################################################################
# .. figure:: arclength_feature.png
#    :align: center
# 
#    Showing the streamlines colored according to their length.
# 
# .. _clustering-examples-VectorOfEndpointsFeature:
# 
# Vector Between Endpoints Feature
# ================================
# **What:** Instances of `VectorOfEndpointsFeature` extract the vector going
# from one extremity of the streamline to the other. In other words, this feature
# represents the vector beginning at the first point and ending at the last point
# of the streamlines.
# 
# **When:** This feature can be useful when you *only* need information about the
# orientation of a streamline.
# 
# **Note:** Since streamlines endpoints are ambiguous (e.g. the first point could
# be either the beginning or the end of the streamline), one must be careful when
# using this feature.


import numpy as np
from dipy.viz import window, colormap
from dipy.segment.clustering import QuickBundles
from dipy.segment.featurespeed import VectorOfEndpointsFeature
from dipy.segment.metric import CosineMetric

# Get some streamlines.
streamlines = get_streamlines()  # Previously defined.

feature = VectorOfEndpointsFeature()
metric = CosineMetric(feature)
qb = QuickBundles(threshold=0.1, metric=metric)
clusters = qb.cluster(streamlines)

# Color each streamline according to the cluster they belong to.
colormap = colormap.create_colormap(np.arange(len(clusters)))
colormap_full = np.ones((len(streamlines), 3))
for cluster, color in zip(clusters, colormap):
    colormap_full[cluster.indices] = color

# Visualization
scene = window.Scene()
scene.clear()
scene.SetBackground(0, 0, 0)
scene.add(actor.streamtube(streamlines, colormap_full))
window.record(scene, out_path='vector_of_endpoints_feature.png',
              size=(600, 600))
if interactive:
    window.show(scene)

###############################################################################
# .. figure:: vector_of_endpoints_feature.png
#    :align: center
# 
#    Showing the streamlines colored according to their orientation.
# 
# .. include:: ../links_names.inc
# 
# .. [Garyfallidis12] Garyfallidis E. et al., QuickBundles a method for
#    tractography simplification, Frontiers in Neuroscience, vol 6, no 175, 2012.