# This script is for reconstructing simulated phantoms generated by TCT_generate_phantom.py

import sys
import mbirjax as mj
import mbirjax.preprocess as mjp
import numpy as np
import os

def main():
    # Recon parameters
    sharpness = 2.0
    num_recon_rows = 40

    # Load the ground truth phantom, synthetic sinogram and reconstruction parameters
    simulated_data_path = "./output/TCT_simulated_data.npz"

    try:
        simulated_data = np.load(simulated_data_path, allow_pickle=True)
    except FileNotFoundError:
        raise FileNotFoundError(
            f"Cannot find '{simulated_data_path}'. "
            "Please run TCT_generate_phantom.py first to generate the simulated data."
        )

    gt_phantom = simulated_data["phantom"]
    sinogram = simulated_data["sino"]
    translation_params = simulated_data["translation_params"].item()
    optional_params = simulated_data["optional_params"].item()

    # Output path
    output_path = './output/'  # path to store output recon images

    # Initialize model for reconstruction.
    tct_model = mj.TranslationModel(**translation_params)
    tct_model.set_params(**optional_params)
    tct_model.set_params(sharpness=sharpness)

    # Set parameters for recon
    sino_shape = sinogram.shape
    source_detector_dist, source_iso_dist = tct_model.get_params(['source_detector_dist', 'source_iso_dist'])
    delta_det_row, delta_det_channel = tct_model.get_params(['delta_det_row', 'delta_det_channel'])
    translation_vectors = tct_model.get_params('translation_vectors')
    recon_shape, _, delta_recon_row = mj.utilities.calc_tct_recon_params(source_detector_dist,
                                                                         source_iso_dist, delta_det_row,
                                                                         delta_det_channel, sino_shape,
                                                                         translation_vectors)
    tct_model.set_params(delta_recon_row=delta_recon_row)
    tct_model.set_params(recon_shape=(num_recon_rows,)+recon_shape[1:])
    tct_model.set_params(positivity_flag=True)
    tct_model.set_params(partition_sequence=5*[0,] + 100*[1, 3,])
    tct_model.set_params(qggmrf_nbr_wts=[0.1, 1, 1])

    # Print model parameters and display translation array
    translation_vectors = translation_params['translation_vectors']
    tct_model.print_params()
    delta_voxel, delta_recon_row = tct_model.get_params(['delta_voxel', 'delta_recon_row'])
    translation_vectors_display = translation_vectors.copy()
    translation_vectors_display[:, 0] /= delta_voxel
    translation_vectors_display[:, 2] /= delta_voxel
    translation_vectors_display[:, 1] /= delta_recon_row
    mj.display_translation_vectors(translation_vectors_display, recon_shape)

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

    # Perform MBIR reconstruction
    direct_recon, direct_dict = tct_model.recon(sinogram, max_iterations=0)
    mbir_recon, mbir_dict = tct_model.recon(sinogram, max_iterations=200, stop_threshold_change_pct=0.1)

    # Save reconstruction results
    os.makedirs(output_path, exist_ok=True)
    output_path = os.path.join(output_path, f'TCT_simulation_recon.h5')
    mj.export_recon_hdf5(output_path, mbir_recon, recon_dict=mbir_dict, top_margin=0, bottom_margin=0)

    # Display results
    mj.slice_viewer(gt_phantom.transpose(0, 2, 1), direct_recon.transpose(0, 2, 1), mbir_recon.transpose(0, 2, 1),
                    data_dicts=[None, direct_dict, mbir_dict],
                    vmin=0, vmax=1.2, title='Object (left), FDK recon (middle), MBIR reconstruction (right)',
                    slice_axis=0)


if __name__ == '__main__':
    main()
