# -*- coding: utf-8 -*-
"""demo_4_wrong_rotation_direction.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1Gd-fMm3XK1WBsuJUklHdZ-4jjsvdpeIT

**MBIRJAX: Wrong Rotation Direction Demo**

See the [MBIRJAX documentation](https://mbirjax.readthedocs.io/en/latest/) for an overview and details.  

This script demonstrates the effects of an incorrect rotation direction when using cone beam projection.  

For this demo, we do a baseline reconstruction with the correct specification of angles, then reverse the order of the angles and perform another reconstruction.  

Since this problem doesn't arise with parallel beam because of symmetry, we show this only for cone beam.   

See [demo_1_shepp_logan.py](https://colab.research.google.com/drive/1zG_H6CDjuQxeMRQHan3XEyX2YVKcSSNC) for the basic steps of synthetic sinogram generation and reconstruction.

Select a GPU as runtime type for best performance.
"""

# Commented out IPython magic to ensure Python compatibility.
# %pip install mbirjax

import numpy as np
import pprint
import time
import jax.numpy as jnp
import mbirjax as mj

"""**Set the geometry parameters**"""

# Set parameters for the problem size.
num_views = 64
num_det_rows = 40
num_det_channels = 128

sharpness = 0.0

# For cone beam geometry, we need to describe the distances source to detector and source to rotation axis.
# np.Inf is an allowable value, in which case this is essentially parallel beam
source_detector_dist = 4 * num_det_channels
source_iso_dist = source_detector_dist

# For cone beam reconstruction, we need a little more than 180 degrees for full coverage.
detector_cone_angle = 2 * np.arctan2(num_det_channels / 2, source_detector_dist)
start_angle = -(np.pi + detector_cone_angle) * (1/2)
end_angle = (np.pi + detector_cone_angle) * (1/2)

"""**Data generation:** For demo purposes, we create a phantom and then project it to create a sinogram.

For this demo we use the default recon shape for cone beam.

Note:  the sliders on the viewer won't work in notebook form.  For that you'll need to run the python code with an interactive matplotlib backend, typcially using the command line or a development environment like Spyder or Pycharm to invoke python.  

"""

# Choose the geometry type - for wrong rotation, only conebeam geometry will have a problem.
model_type = 'cone'
object_type = 'shepp-logan'  # 'shepp-logan' or 'cube'


# Generate simulated data
# In a real application you would not have the phantom, but we include it here for later display purposes
phantom, sinogram, params = mj.generate_demo_data(object_type=object_type, model_type=model_type,
                                                  num_views=num_views, num_det_rows=num_det_rows,
                                                  num_det_channels=num_det_channels)
angles = params['angles']


# View sinogram
title='Original sinogram'
mj.slice_viewer(sinogram, title=title, slice_axis=0, slice_label='View')

"""**Do a baseline reconstruction**

Here we do a standard reconstruction with the correct specification of angles.  
"""

print('\nStarting recon with correct rotation\n')
ct_model = mj.ConeBeamModel(sinogram.shape, angles, source_detector_dist=source_detector_dist,
                                            source_iso_dist=source_iso_dist)
ct_model.set_params(sharpness=sharpness)

recon_correct, recon_dict_correct = ct_model.recon(sinogram)

recon_params_correct = recon_dict_correct['recon_params']
pprint.pprint(recon_params_correct, compact=True)

"""**Do a reconstruction with angles in reverse order.**"""

print('\nStarting recon with incorrect rotation\n')
angles_reversed = angles[::-1]
ct_model = mj.ConeBeamModel(sinogram.shape, angles_reversed, source_detector_dist=source_detector_dist,
                                            source_iso_dist=source_iso_dist)
recon_incorrect, recon_dict_incorrect = ct_model.recon(sinogram)

recon_params_incorrect = recon_dict_incorrect['recon_params']
pprint.pprint(recon_params_incorrect, compact=True)

"""**Display the results with correct and incorrect angles.**"""

title = 'Correct rotation recon (left) vs incorrect rotation recon (right)'
title += '\nThe incorrect angle specification leads to shape distortion and top/bottom reflection.'
mj.slice_viewer(recon_correct, recon_incorrect, data_dicts=[recon_dict_correct, recon_dict_incorrect], title=title, vmin=0.0, vmax=1.0)

"""**Next:** Try changing some of the parameters and re-running or try [some of the other demos](https://mbirjax.readthedocs.io/en/latest/demos_and_faqs.html).  """