"""
======================
Fast Streamline Search
======================

This example explains how Fast Streamline Search [StOnge2022]_
can be used to find all similar streamlines.

First import the necessary modules.
"""

import numpy as np

from dipy.data import (get_target_tractogram_hcp, get_two_hcp842_bundles,
                       fetch_bundle_atlas_hcp842, fetch_target_tractogram_hcp)
from dipy.io.streamline import load_trk
from dipy.segment.fss import FastStreamlineSearch, nearest_from_matrix_row
from dipy.viz import actor, window

"""
Download and read data for this tutorial
"""
fetch_bundle_atlas_hcp842()
fetch_target_tractogram_hcp()

hcp_file = get_target_tractogram_hcp()
streamlines = load_trk(hcp_file, "same", bbox_valid_check=False).streamlines

"""
Visualize the atlas (ref) bundle and full brain tractogram
"""

interactive = False

scene = window.Scene()
scene.SetBackground(1, 1, 1)
scene.add(actor.line(streamlines))
if interactive:
    window.show(scene)
else:
    window.record(scene, out_path='tractograms_initial.png', size=(600, 600))

"""
.. figure:: tractograms_initial.png
   :align: center

   Atlas bundle and source streamlines before registration.

"""

"""
Read Arcuate Fasciculus Left and Corticospinal Tract Left bundles from already
fetched atlas data to use them as model bundle. Let's visualize the
Arcuate Fasciculus Left model bundle.
"""

model_af_l_file, model_cst_l_file = get_two_hcp842_bundles()
sft_af_l = load_trk(model_af_l_file, "same", bbox_valid_check=False)
model_af_l = sft_af_l.streamlines

scene = window.Scene()
scene.SetBackground(1, 1, 1)
scene.add(actor.line(model_af_l, colors=(0, 1, 0)))
scene.set_camera(focal_point=(-18.17281532, -19.55606842, 6.92485857),
                 position=(-360.11, -30.46, -40.44),
                 view_up=(-0.03, 0.028, 0.89))
if interactive:
    window.show(scene)
else:
    window.record(scene, out_path='AF_L_model_bundle.png',
                  size=(600, 600))

"""
.. figure:: AF_L_model_bundle.png
   :align: center

   Model Arcuate Fasciculus Left bundle

"""

"""
Search for all similar streamlines  [StOnge2022]_

Fast Streamline Search can do a radius search
to find all streamlines that are similar to from one tractogram to another.
It return the distance matrix mapping between both tractograms.
The same list of streamlines can be given to recover the self distance matrix.

    FastStreamlinesSearch Class
    - ref_streamlines : reference streamlines, that will be searched in (tree)
    - max_radius : is the maximum distance that can be used with radius search

    radius_search() function
    - radius : for each streamline search find all similar ones in the
        "ref_streamlines" that are within the given radius

    Be cautious, a large radius might result in a dense distance computation,
    requiring a large amount of time and memory.
    Recommended range of the radius is from 1 to 10 mm.
"""

radius = 7.0
fs_tree_af = FastStreamlineSearch(ref_streamlines=model_af_l, max_radius=radius)
coo_mdist_mtx = fs_tree_af.radius_search(streamlines, radius=radius)


"""
Extract indices of streamlines with an similar ones in the reference
"""
ids_s = np.unique(coo_mdist_mtx.row)
ids_ref = np.unique(coo_mdist_mtx.col)

recognized_af_l = streamlines[ids_s]

"""
let's visualize streamlines similar to the Arcuate Fasciculus Left bundle
"""

scene = window.Scene()
scene.SetBackground(1, 1, 1)
scene.add(actor.line(recognized_af_l, colors=(0, 0, 1)))
scene.set_camera(focal_point=(-18.17281532, -19.55606842, 6.92485857),
                 position=(-360.11, -30.46, -40.44),
                 view_up=(-0.03, 0.028, 0.89))
if interactive:
    window.show(scene)
else:
    window.record(scene, out_path='AF_L_recognized_bundle.png',
                  size=(600, 600))

"""
.. figure:: AF_L_recognized_bundle.png
   :align: center

   Recognized Arcuate Fasciculus Left bundle

"""

"""
Color the resulting AF by the nearest streamlines distance
"""

nn_s, nn_ref, nn_dist = nearest_from_matrix_row(coo_mdist_mtx)

scene = window.Scene()
scene.SetBackground(1, 1, 1)
cmap = actor.colormap_lookup_table(scale_range=(nn_dist.min(), nn_dist.max()))
scene.add(actor.line(recognized_af_l, colors=nn_dist, lookup_colormap=cmap))
scene.add(actor.scalar_bar(cmap, title="distance to atlas (mm)"))
scene.set_camera(focal_point=(-18.17281532, -19.55606842, 6.92485857),
                 position=(-360.11, -30.46, -40.44),
                 view_up=(-0.03, 0.028, 0.89))
if interactive:
    window.show(scene)
else:
    window.record(scene, out_path='AF_L_recognized_bundle_dist.png',
                  size=(600, 600))

"""
.. figure:: AF_L_recognized_bundle_dist.png
   :align: center

   Recognized Arcuate Fasciculus Left bundle colored by distance to ref

"""

"""
Display the streamlines reference reached in green
"""
# Default red color
ref_color = np.zeros((len(model_af_l), 3), dtype=float)
ref_color[:] = (1.0, 0.0, 0.0)

# Reached in green color
ref_color[ids_ref] = (0.0, 1.0, 0.0)

scene = window.Scene()
scene.SetBackground(1, 1, 1)
scene.add(actor.line(model_af_l, ref_color))
scene.set_camera(focal_point=(-18.17281532, -19.55606842, 6.92485857),
                 position=(-360.11, -30.46, -40.44),
                 view_up=(-0.03, 0.028, 0.89))

if interactive:
    window.show(scene)
else:
    window.record(scene, out_path='AF_L_model_bundle_reached.png',
                  size=(600, 600))

"""
.. figure:: AF_L_model_bundle_reached.png
   :align: center

   Arcuate Fasciculus Left model reached (green) in radius

"""

"""
References
----------

.. [StOnge2022] St-Onge E. et al. Fast Streamline Search:
                An Exact Technique for Diffusion MRI Tractography.
                Neuroinformatics, 2022.
"""
