# Imports
from mri.operators import FFT, WaveletN
from mri.operators.utils import convert_mask_to_locations
from mri.reconstructors import SingleChannelReconstructor
from mri.scripts.gridsearch import launch_grid

from pysap.data import get_sample_data

from modopt.math.metrics import ssim
from modopt.opt.proximity import SparseThreshold
from modopt.opt.linear import Identity
import numpy as np


# Load MR data and obtain kspace
image = get_sample_data('2d-mri')
mask = get_sample_data("cartesian-mri-mask")
kspace_loc = convert_mask_to_locations(mask.data)
fourier_op = FFT(samples=kspace_loc, shape=image.shape)
kspace_data = fourier_op.op(image.data)
# Define the keyword dictionaries based on convention
ref = image
metrics = {
    'ssim': {
        'metric': ssim,
        'mapping': {'x_new': 'test', 'y_new': None},
        'cst_kwargs': {'ref': image, 'mask': None},
        'early_stopping': True,
    },
}
linear_params = {
    'init_class': WaveletN,
    'kwargs':
        {
            'wavelet_name': ['sym8', 'sym12'],
            'nb_scale': [3, 4]
        }
}
regularizer_params = {
    'init_class': SparseThreshold,
    'kwargs':
        {
            'linear': Identity(),
            'weights': np.logspace(-8, -6, 5),
        }
}
optimizer_params = {
    # Just following convention
    'kwargs':
        {
            'optimization_alg': 'fista',
            'num_iterations': 20,
            'metrics': metrics,
        }
}
# Call the launch grid function and obtain results
raw_results, test_cases, key_names, best_idx = launch_grid(
    kspace_data=kspace_data,
    fourier_op=fourier_op,
    linear_params=linear_params,
    regularizer_params=regularizer_params,
    optimizer_params=optimizer_params,
    reconstructor_class=SingleChannelReconstructor,
    reconstructor_kwargs={'gradient_formulation': 'synthesis'},
    compare_metric_details={'metric': 'ssim'},
    n_jobs=-1,
    verbose=1,
)
image_rec = raw_results[best_idx][0]