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

Automatically generated by Colab.

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

**MBIRJAX: Basic Demo**

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

This script demonstrates the basic MBIRJAX code by creating a 3D phantom inspired by Shepp-Logan, forward projecting it to create a sinogram, and then using MBIRJAX to perform a Model-Based, Multi-Granular Vectorized Coordinate Descent reconstruction.

For the demo, we create some synthetic data by first making a phantom, then forward projecting it to obtain a sinogram.  

In a real application, you would load your sinogram as a numpy array and use numpy.transpose if needed so that it
has axes in the order (views, rows, channels).  For reference, assuming the rotation axis is vertical, then increasing the row index nominally moves down the rotation axis and increasing the channel index moves to the right as seen from the source.

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

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

import numpy as np
import time
import pprint
import jax.numpy as jnp
import mbirjax as mj
import matplotlib.pyplot as plt

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

def main():
    # Choose the geometry type
    geometry_type = 'parallel'  # 'cone' or 'parallel'

    # Set parameters for the problem size - you can vary these, but if you make num_det_rows very small relative to
    # channels, then the generated phantom may not have an interior.
    num_views = 32
    num_det_rows = 60
    num_det_channels = 256
    sharpness = 4.0
    snr_db = 40

    # 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.
    if geometry_type == 'cone':
        detector_cone_angle = 2 * np.arctan2(num_det_channels / 2, source_detector_dist)
    else:
        detector_cone_angle = 0
    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.
    
    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.  
    
    """

    # Initialize sinogram
    sinogram_shape = (num_views, num_det_rows, num_det_channels)
    angles = jnp.linspace(start_angle, end_angle, num_views, endpoint=False)

    if geometry_type == 'cone':
        ct_model_for_generation = mj.ConeBeamModel(sinogram_shape, angles, source_detector_dist=source_detector_dist, source_iso_dist=source_iso_dist)
    elif geometry_type == 'parallel':
        ct_model_for_generation = mj.ParallelBeamModel(sinogram_shape, angles)
    else:
        raise ValueError('Invalid geometry type.  Expected cone or parallel, got {}'.format(geometry_type))

    # Generate 3D Shepp Logan phantom
    print('Creating phantom')
    phantom_shape = ct_model_for_generation.get_params('recon_shape')
    phantom = mj.generate_3d_shepp_logan_low_dynamic_range(phantom_shape)

    # Generate synthetic sinogram data
    print('Creating sinogram')
    sinogram = ct_model_for_generation.forward_project(phantom)
    sinogram = np.array(sinogram)

    # View sinogram
    title = 'Original sinogram \nUse the sliders to change the view or adjust the intensity range.'
    # mj.slice_viewer(sinogram, slice_axis=0, title=title, slice_label='View')

    """**Initialize for the reconstruction**"""

    # ####################
    # Initialize the model for reconstruction.
    if geometry_type == 'cone':
        ct_model_for_recon = mj.ConeBeamModel(sinogram_shape, angles, source_detector_dist=source_detector_dist, source_iso_dist=source_iso_dist)
    else:
        ct_model_for_recon = mj.ParallelBeamModel(sinogram_shape, angles)

    # Generate weights array - for an initial reconstruction, use weights = None, then modify if needed.
    weights = None
    # weights = mj.gen_weights(sinogram / sinogram.max(), weight_type='transmission_root')

    # Set parameters
    partition_sequence = ct_model_for_recon.get_params('partition_sequence')
    ct_model_for_recon.set_params(sharpness=sharpness, snr_db=snr_db, partition_sequence=partition_sequence)

    # Print out model parameters
    print('Evaluating convergence with sharpness = {}, snr_db = {}'.format(sharpness, snr_db))

    # Get the baseline
    compute_baseline = False
    output_dir = '../convergence_output'
    os.makedirs(output_dir, exist_ok=True)

    max_iterations = 1000
    iterations_per_step = 5
    stop_threshold_change_pct = 0.0

    short_iterations = 30

    if compute_baseline:

        baseline, nrmse_baseline = recon_by_continuation(ct_model_for_recon, sinogram, weights,
                                                         max_iterations, iterations_per_step, stop_threshold_change_pct,
                                                         baseline=None)
        np.savez(os.path.join(output_dir, 'baseline.npz'), baseline)

        # Repeat to show convergence
        recon = None
        default_nrmse = np.zeros(np.ceil(max_iterations / iterations_per_step).astype(int))

        for iteration in range(0, max_iterations, iterations_per_step):
            i = iteration // iterations_per_step
            recon, recon_dict = ct_model_for_recon.recon(sinogram, weights=weights,
                                                           max_iterations=iteration + iterations_per_step,
                                                           first_iteration=iteration, init_recon=recon,
                                                           compute_prior_loss=True,
                                                           stop_threshold_change_pct=stop_threshold_change_pct)
            default_nrmse[i] = np.linalg.norm(recon - baseline) / np.linalg.norm(baseline)

        np.savez(os.path.join(output_dir, 'default.npz'), recon)
        np.savez(os.path.join(output_dir, 'default_nrmse.npz'), default_nrmse)

        baseline_short, baseline_short_nrmse = recon_by_continuation(ct_model_for_recon, sinogram, weights,
                                                                     short_iterations, iterations_per_step,
                                                                     stop_threshold_change_pct,
                                                                     baseline=baseline)
        np.savez(os.path.join(output_dir, 'baseline_short.npz'), baseline_short)
        np.savez(os.path.join(output_dir, 'baseline_short_nrmse.npz'), baseline_short_nrmse)

    try:
        baseline = np.load(os.path.join(output_dir, 'baseline.npz'))['arr_0']
        default_nrmse = np.load(os.path.join(output_dir, 'default_nrmse.npz'))['arr_0']
        baseline_short_nrmse = np.load(os.path.join(output_dir, 'baseline_short_nrmse.npz'))['arr_0']

        plt.figure(1)
        plt.plot(iterations_per_step * np.arange(len(default_nrmse)), default_nrmse)
        plt.plot(iterations_per_step * np.arange(len(baseline_short_nrmse)), baseline_short_nrmse)
        # plt.axis((0, 30, 0.01, 0.08))
        plt.title('NRMSE relative to baseline vs iteration\nfor default recon and recon with increasing sharpness and snr_db')
        plt.legend(['Default', 'New varying sharpness/snr_db'])
        plt.show()

    except FileNotFoundError as e:
        raise FileNotFoundError('baseline and/or default files not found.  Rerun with compute_baseline = True')
    # recon_1000_nrmse = np.load(os.path.join(output_dir, 'default_nrmse.npz'))['arr_0']
    # plt.plot(iterations_per_step * np.arange(len(default_nrmse)), default_nrmse)
    # plt.plot(iterations_per_step * np.arange(len(recon_1000_nrmse)), recon_1000_nrmse)
    # plt.axis((0, 30, 0.01, 0.08))
    # plt.legend(['Convergence to baseline', 'Convergence with varying sharpness/snr_db'])

    # ##########################
    # Perform reconstruction with varying sharpness/snr_db
    recon, recon_nrmse = recon_by_continuation(ct_model_for_recon, sinogram, weights,
                                               short_iterations, iterations_per_step, stop_threshold_change_pct,
                                               baseline=baseline)

    np.savez(os.path.join(output_dir, 'recon.npz'), recon)
    np.savez(os.path.join(output_dir, 'recon_nrmse.npz'), recon_nrmse)

    plt.figure(2)
    plt.plot(iterations_per_step * np.arange(len(default_nrmse)), default_nrmse)
    plt.plot(iterations_per_step * np.arange(len(recon_nrmse)), recon_nrmse)
    plt.axis((0, max_iterations, 0.01, 0.08))
    plt.legend(['Default', 'New varying sharpness/snr_db'])
    plt.show(block=True)
    # plt.axis((0, iterations_per_step * len(default_nrmse), 0, 0.3))
    mj.slice_viewer(recon, recon - baseline, title='New recon (left) and (new recon) - baseline (right)')


def recon_by_continuation(ct_model_for_recon, sinogram, weights,
                          max_iterations, iterations_per_step, stop_threshold_change_pct, baseline=None):
    # ##########################
    # Perform reconstruction with varying sharpness/snr_db
    print('Starting recon')
    recon = None

    sharpness, snr_db, partition_sequence = ct_model_for_recon.get_params(['sharpness', 'snr_db', 'partition_sequence'])
    nrmse = np.zeros(max_iterations // iterations_per_step)
    time0 = time.time()
    for iteration in range(0, max_iterations, iterations_per_step):
        i = iteration // iterations_per_step
        if iteration == 0:
            cur_sharpness = 1
            cur_snr_db = 25
            cur_partition_sequence = [0, 2, 4, 6, 7]
        elif iteration == iterations_per_step:
            cur_sharpness = (cur_sharpness + 3 * sharpness) / 4
            # cur_snr_db = (cur_snr_db + snr_db) / 2
            cur_partition_sequence = [7]
        elif iteration == 2 * iterations_per_step:
            cur_sharpness = sharpness
            # cur_snr_db = (cur_snr_db + snr_db) / 2
            cur_partition_sequence = 2 * iterations_per_step * [0] + [7]
        elif iteration == 3 * iterations_per_step:
            cur_sharpness = sharpness
            cur_snr_db = (cur_snr_db + snr_db) / 2
            cur_partition_sequence = 3 * iterations_per_step * [0] + [4, 6, 7]
        else:
            cur_sharpness = sharpness  # if iteration >= iterations_per_step else 1
            cur_snr_db = snr_db  # if iteration >= iterations_per_step else 30
            cur_partition_sequence = 4 * iterations_per_step * [0] + [4, 6, 7]
        ct_model_for_recon.set_params(sharpness=cur_sharpness, snr_db=cur_snr_db, partition_sequence=cur_partition_sequence)
        recon, recon_dict = ct_model_for_recon.recon(sinogram, weights=weights,
                                                       max_iterations=iteration + iterations_per_step,
                                                       first_iteration=iteration, init_recon=recon,
                                                       compute_prior_loss=True,
                                                       stop_threshold_change_pct=stop_threshold_change_pct)
        if baseline is not None:
            nrmse[i] = np.linalg.norm(recon - baseline) / np.linalg.norm(baseline)

        recon_params = recon_dict['recon_params']
        if recon_params['stop_threshold_change_pct'][-1] < stop_threshold_change_pct:
            break

    elapsed = time.time() - time0
    print('Elapsed time for recon is {:.3f} seconds'.format(elapsed))

    ct_model_for_recon.set_params(sharpness=sharpness, snr_db=snr_db, partition_sequence=partition_sequence)
    return recon, nrmse


if __name__ == "__main__":
    main()