"""
========================================
Visualize bundles and metrics on bundles
========================================

First, let's download some available datasets. Here we are using a dataset
which provides metrics and bundles.
"""

import numpy as np
from dipy.viz import window, actor
from dipy.data import fetch_bundles_2_subjects, read_bundles_2_subjects
from dipy.tracking.streamline import transform_streamlines

fetch_bundles_2_subjects()
dix = read_bundles_2_subjects(subj_id='subj_1', metrics=['fa'],
                              bundles=['cg.left', 'cst.right'])

"""
Store fractional anisotropy.
"""

fa = dix['fa']

"""
Store grid to world transformation matrix.
"""

affine = dix['affine']

"""
Store the cingulum bundle. A bundle is a list of streamlines.
"""

bundle = dix['cg.left']

"""
It happened that this bundle is in world coordinates and therefore we need to
transform it into native image coordinates so that it is in the same coordinate
space as the ``fa`` image.
"""

bundle_native = transform_streamlines(bundle, np.linalg.inv(affine))

"""
Show every streamline with an orientation color
===============================================

This is the default option when you are using ``line`` or ``streamtube``.
"""

scene = window.Scene()

stream_actor = actor.line(bundle_native)

scene.set_camera(position=(-176.42, 118.52, 128.20),
                 focal_point=(113.30, 128.31, 76.56),
                 view_up=(0.18, 0.00, 0.98))

scene.add(stream_actor)

# Uncomment the line below to show to display the window
# window.show(scene, size=(600, 600), reset_camera=False)
window.record(scene, out_path='bundle1.png', size=(600, 600))

"""
.. figure:: bundle1.png
   :align: center

   One orientation color for every streamline.

You may wonder how we knew how to set the camera. This is very easy. You just
need to run ``window.show`` once to see how you want to see the object and then
close the window and call the ``camera_info`` method which prints the position,
focal point and view up vectors of the camera.
"""

scene.camera_info()

"""
Show every point with a value from a volume with default colormap
=================================================================

Here we will need to input the ``fa`` map in ``streamtube`` or ``line``.
"""

scene.clear()
stream_actor2 = actor.line(bundle_native, fa, linewidth=0.1)

"""
We can also show the scalar bar.
"""

bar = actor.scalar_bar()

scene.add(stream_actor2)
scene.add(bar)

# window.show(scene, size=(600, 600), reset_camera=False)
window.record(scene, out_path='bundle2.png', size=(600, 600))

"""
.. figure:: bundle2.png
   :align: center

   Every point with a color from FA.

Show every point with a value from a volume with your colormap
==============================================================

Here we will need to input the ``fa`` map in ``streamtube``
"""

scene.clear()

hue = (0.0, 0.0)  # red only
saturation = (0.0, 1.0)  # white to red

lut_cmap = actor.colormap_lookup_table(hue_range=hue,
                                       saturation_range=saturation)

stream_actor3 = actor.line(bundle_native, fa, linewidth=0.1,
                           lookup_colormap=lut_cmap)
bar2 = actor.scalar_bar(lut_cmap)

scene.add(stream_actor3)
scene.add(bar2)

# window.show(scene, size=(600, 600), reset_camera=False)
window.record(scene, out_path='bundle3.png', size=(600, 600))

"""
.. figure:: bundle3.png
   :align: center

   Every point with a color from FA using a non default colormap.


Show every bundle with a specific color
========================================

You can have a bundle with a specific color. In this example, we are choosing
orange.
"""

scene.clear()
stream_actor4 = actor.line(bundle_native, (1., 0.5, 0), linewidth=0.1)

scene.add(stream_actor4)

# window.show(scene, size=(600, 600), reset_camera=False)
window.record(scene, out_path='bundle4.png', size=(600, 600))

"""
.. figure:: bundle4.png
   :align: center

   Entire bundle with a specific color.

Show every streamline of a bundle with a different color
========================================================

Let's make a colormap where every streamline of the bundle is colored by its
length.
"""

scene.clear()

from dipy.tracking.streamline import length

lengths = length(bundle_native)

hue = (0.5, 0.5)  # blue only
saturation = (0.0, 1.0)  # black to white

lut_cmap = actor.colormap_lookup_table(
    scale_range=(lengths.min(), lengths.max()),
    hue_range=hue,
    saturation_range=saturation)

stream_actor5 = actor.line(bundle_native, lengths, linewidth=0.1,
                           lookup_colormap=lut_cmap)

scene.add(stream_actor5)
bar3 = actor.scalar_bar(lut_cmap)

scene.add(bar3)

# window.show(scene, size=(600, 600), reset_camera=False)
window.record(scene, out_path='bundle5.png', size=(600, 600))

"""
.. figure:: bundle5.png
   :align: center

   Color every streamline by the length of the streamline


Show every point of every streamline with a different color
============================================================

In this case in which we want to have a color per point and per streamline,
we can create a list of the colors to correspond to the list of streamlines
(bundles). Here in ``colors`` we will insert some random RGB colors.
"""

scene.clear()

colors = [np.random.rand(*streamline.shape) for streamline in bundle_native]

stream_actor6 = actor.line(bundle_native, np.vstack(colors), linewidth=0.2)

scene.add(stream_actor6)

# window.show(scene, size=(600, 600), reset_camera=False)
window.record(scene, out_path='bundle6.png', size=(600, 600))

"""
.. figure:: bundle6.png
   :align: center

   Random colors per point per streamline.

In summary, we showed that there are many useful ways for visualizing maps
on bundles.

"""