
"""
=====================
Simple volume slicing
=====================

Here we present an example for visualizing slices from 3D images.

"""

import os

from dipy.data import fetch_bundles_2_subjects
from dipy.io.image import load_nifti, load_nifti_data
from dipy.viz import window, actor, ui

"""
Let's download and load a T1.
"""

fetch_bundles_2_subjects()

fname_t1 = os.path.join(os.path.expanduser('~'), '.dipy',
                        'exp_bundles_and_maps', 'bundles_2_subjects',
                        'subj_1', 't1_warped.nii.gz')

data, affine = load_nifti(fname_t1)

"""
Create a Scene object which holds all the actors which we want to visualize.
"""

scene = window.Scene()
scene.background((0.5, 0.5, 0.5))

"""
Render slices from T1 with a specific value range
=================================================

The T1 has usually a higher range of values than what can be visualized in an
image. We can set the range that we would like to see.
"""

mean, std = data[data > 0].mean(), data[data > 0].std()
value_range = (mean - 0.5 * std, mean + 1.5 * std)

"""
The ``slice`` function will read data and resample the data using an affine
transformation matrix. The default behavior of this function is to show the
middle slice of the last dimension of the resampled data.
"""

slice_actor = actor.slicer(data, affine, value_range)

"""
The ``slice_actor`` contains an axial slice.
"""

scene.add(slice_actor)

"""
The same actor can show any different slice from the given data using its
``display`` function. However, if we want to show multiple slices we need to
copy the actor first.
"""

slice_actor2 = slice_actor.copy()

"""
Now we have a new ``slice_actor`` which displays the middle slice of the
sagittal plane.
"""

slice_actor2.display(slice_actor2.shape[0]//2, None, None)

scene.add(slice_actor2)

scene.reset_camera()
scene.zoom(1.4)

"""
In order to interact with the data you will need to uncomment the line below.
"""

# window.show(scene, size=(600, 600), reset_camera=False)

"""
Otherwise, you can save a screenshot using the following command.
"""

window.record(scene, out_path='slices.png', size=(600, 600),
              reset_camera=False)

"""
.. figure:: slices.png
   :align: center

   Simple slice viewer.

Render slices from FA with your colormap
========================================

It is also possible to set the colormap of your preference. Here we are loading
an FA image and showing it in a non-standard way using an HSV colormap.
"""

fname_fa = os.path.join(os.path.expanduser('~'), '.dipy',
                        'exp_bundles_and_maps', 'bundles_2_subjects',
                        'subj_1', 'fa_1x1x1.nii.gz')

fa = load_nifti_data(fname_fa)

lut = actor.colormap_lookup_table(scale_range=(fa.min(), fa.max()*0.8),
                                  hue_range=(0.4, 1.),
                                  saturation_range=(1, 1.),
                                  value_range=(0., 1.))

"""
This is because the lookup table is applied in the slice after interpolating
to (0, 255).
"""

fa_actor = actor.slicer(fa, affine, lookup_colormap=lut)

scene.clear()
scene.add(fa_actor)

scene.reset_camera()
scene.zoom(1.4)

# window.show(scene, size=(600, 600), reset_camera=False)

window.record(scene, out_path='slices_lut.png', size=(600, 600),
              reset_camera=False)

"""
.. figure:: slices_lut.png
   :align: center

   **Simple slice viewer with an HSV colormap**.
"""

"""
Now we would like to add the ability to click on a voxel and show its value
on a panel in the window. The panel is a UI element which requires access to
different areas of the visualization pipeline and therefore we don't recommend
using it with ``window.show``. The more appropriate way is to use the
``ShowManager`` object, which allows accessing the pipeline in different areas.
"""

show_m = window.ShowManager(scene, size=(1200, 900))
show_m.initialize()

"""
We'll start by creating the panel and adding it to the ``ShowManager``
"""

label_position = ui.TextBlock2D(text='Position:')
label_value = ui.TextBlock2D(text='Value:')

result_position = ui.TextBlock2D(text='')
result_value = ui.TextBlock2D(text='')

panel_picking = ui.Panel2D(size=(250, 125),
                           position=(20, 20),
                           color=(0, 0, 0),
                           opacity=0.75,
                           align="left")

panel_picking.add_element(label_position, (0.1, 0.55))
panel_picking.add_element(label_value, (0.1, 0.25))

panel_picking.add_element(result_position, (0.45, 0.55))
panel_picking.add_element(result_value, (0.45, 0.25))

scene.add(panel_picking)

"""
Add a left-click callback to the slicer. Also disable interpolation so you can
see what you are picking.
"""


def left_click_callback(obj, ev):
    """Get the value of the clicked voxel and show it in the panel."""
    event_pos = show_m.iren.GetEventPosition()

    obj.picker.Pick(event_pos[0], event_pos[1], 0, scene)

    i, j, k = obj.picker.GetPointIJK()
    result_position.message = '({}, {}, {})'.format(str(i), str(j), str(k))
    result_value.message = '%.8f' % data[i, j, k]


fa_actor.SetInterpolate(False)
fa_actor.AddObserver('LeftButtonPressEvent', left_click_callback, 1.0)

# show_m.start()

"""
Create a mosaic
===============

By using the ``copy`` and ``display`` method of the ``slice_actor`` it becomes
easy and efficient to create a mosaic of all the slices.

So, let's clear the scene and change the projection from perspective to
parallel. We'll also need a new show manager and an associated callback.
"""

scene.clear()
scene.projection('parallel')

result_position.message = ''
result_value.message = ''

show_m_mosaic = window.ShowManager(scene, size=(1200, 900))
show_m_mosaic.initialize()


def left_click_callback_mosaic(obj, ev):
    """Get the value of the clicked voxel and show it in the panel."""
    event_pos = show_m_mosaic.iren.GetEventPosition()

    obj.picker.Pick(event_pos[0], event_pos[1], 0, scene)

    i, j, k = obj.picker.GetPointIJK()
    result_position.message = '({}, {}, {})'.format(str(i), str(j), str(k))
    result_value.message = '%.8f' % data[i, j, k]


"""
Now we need to create two nested for loops which will set the positions of
the grid of the mosaic and add the new actors to the scene. We are going
to use 15 columns and 10 rows but you can adjust those with your datasets.
"""

cnt = 0

X, Y, Z = slice_actor.shape[:3]

rows = 10
cols = 15
border = 10

for j in range(rows):
    for i in range(cols):
        slice_mosaic = slice_actor.copy()
        slice_mosaic.display(None, None, cnt)
        slice_mosaic.SetPosition((X + border) * i,
                                 0.5 * cols * (Y + border) - (Y + border) * j,
                                 0)
        slice_mosaic.SetInterpolate(False)
        slice_mosaic.AddObserver('LeftButtonPressEvent',
                                 left_click_callback_mosaic,
                                 1.0)
        scene.add(slice_mosaic)
        cnt += 1
        if cnt > Z:
            break
    if cnt > Z:
        break

scene.reset_camera()
scene.zoom(1.0)

# show_m_mosaic.ren.add(panel_picking)
# show_m_mosaic.start()

"""
If you uncomment the two lines above, you will be able to move
the mosaic up/down and left/right using the middle mouse button drag,
zoom in/out using the scroll wheel, and pick voxels with left click.
"""

window.record(scene, out_path='mosaic.png', size=(900, 600),
              reset_camera=False)

"""
.. figure:: mosaic.png
   :align: center

   A mosaic of all the slices in the T1 volume.
"""
