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

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1WQwIJ_mDcuMMcWseM66aRPvtv6FmMWF-

**MBIRJAX: Cropped Center Reconstruction Demo**

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

This script demonstrates how to do a reconstruction restricted to an area around the center of rotation.  This is a difficult problem for model-based methods, so the reconstruction has some unavoidable artifacts and intensity shifts, but it does show features near the center without having to do a full reconstruction, which may be useful for large sinograms.  

For this demo, we use a larger detector since the artifacts are more pronounced for very small-scale problems.  To maintain a short run-time, we reduce the number of slices.

Also, we use a phantom that projects fully onto the detector, but you can use the same method when the object projects partially outside the detector.  

For simplicity, we show this only for parallel beam, but the same steps apply 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.  Here we use a fairly large number of views and channels to illustrate
# the behavior of a center cropped recon.
num_views = 400
num_det_rows = 20
num_det_channels = 400

"""**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 parallel 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
model_type = 'parallel'  # 'cone' or 'parallel'
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\nVery few rows to allow for fast execution'
mj.slice_viewer(sinogram, title=title, slice_axis=0, slice_label='View')

"""**Do a cropped center reconstruction**

We specify a recon size smaller than the default. This specifies that only a subregion near the center of rotation will be reconstructed.  

As in demo_2_large_object, this will yield artifacts because some of the information in the sinogram comes from voxels that are projected to the detector on only some of the views.  We minimize the high-frequency artifacts by decreasing sharpness a little, but thee is a residual intensity shift.  
"""

# Initialize model for reconstruction.
weights = None
# ####################
# Use the parameters to get the data and initialize the model for reconstruction.
if model_type == 'cone':
    source_detector_dist = params['source_detector_dist']
    source_iso_dist = params['source_iso_dist']
    ct_model = mj.ConeBeamModel(sinogram.shape, angles, source_detector_dist=source_detector_dist, source_iso_dist=source_iso_dist)
else:
    ct_model = mj.ParallelBeamModel(sinogram.shape, angles)

# Print model parameters
ct_model.print_params()

sharpness = -0.5
recon_row_scale = 0.5
recon_col_scale = 0.5
ct_model.scale_recon_shape(row_scale=recon_row_scale, col_scale=recon_col_scale)

ct_model.set_params(sharpness=sharpness)

print('Starting cropped center recon')
time0 = time.time()
recon, recon_dict = ct_model.recon(sinogram, weights=weights)
recon.block_until_ready()
elapsed = time.time() - time0

# Print out parameters used in recon
recon_params = recon_dict['recon_params']
pprint.pprint(recon_params, compact=True)
print('Elapsed time for recon is {:.3f} seconds'.format(elapsed))

"""**Display the cropped center reconstruction.**"""

title = 'Cropped center recon with sharpness = {:.1f}: Phantom (left) vs VCD Recon (right)'.format(sharpness)
title += '\nThis recon does not include all pixels used to generate the sinogram.'
title += '\nThe missing pixels lead to an intensity shift (adjust intensity to [0, 1]) and a bright outer ring.'
recon_shape = ct_model.get_params('recon_shape')
recon_radius = [length // 2 for length in recon_shape]
start_inds = [phantom.shape[j] // 2 - recon_radius[j] for j in range(2)]
end_inds = [start_inds[j] + recon_shape[j] for j in range(2)]
cropped_phantom = phantom[start_inds[0]:end_inds[0], start_inds[1]:end_inds[1]]
mj.slice_viewer(cropped_phantom, recon, data_dicts=[None, recon_dict], title=title, vmin=0.0, vmax=2.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).  """