import numpy as np
from scipy import sparse
from scipy import ndimage
from sklearn.linear_model import Lasso
from sklearn.linear_model import Ridge
import matplotlib.pyplot as plt
import gradio as gr


def _weights(x, dx=1, orig=0):
    x = np.ravel(x)
    floor_x = np.floor((x - orig) / dx).astype(np.int64)
    alpha = (x - orig - floor_x * dx) / dx
    return np.hstack((floor_x, floor_x + 1)), np.hstack((1 - alpha, alpha))


def _generate_center_coordinates(l_x):
    X, Y = np.mgrid[:l_x, :l_x].astype(np.float64)
    center = l_x / 2.0
    X += 0.5 - center
    Y += 0.5 - center
    return X, Y


def build_projection_operator(l_x, n_dir):
    """Compute the tomography design matrix.
    Parameters
    ----------
    l_x : int
        linear size of image array
    n_dir : int
        number of angles at which projections are acquired.
    Returns
    -------
    p : sparse matrix of shape (n_dir l_x, l_x**2)
    """
    X, Y = _generate_center_coordinates(l_x)
    angles = np.linspace(0, np.pi, n_dir, endpoint=False)
    data_inds, weights, camera_inds = [], [], []
    data_unravel_indices = np.arange(l_x**2)
    data_unravel_indices = np.hstack((data_unravel_indices, data_unravel_indices))
    for i, angle in enumerate(angles):
        Xrot = np.cos(angle) * X - np.sin(angle) * Y
        inds, w = _weights(Xrot, dx=1, orig=X.min())
        mask = np.logical_and(inds >= 0, inds < l_x)
        weights += list(w[mask])
        camera_inds += list(inds[mask] + i * l_x)
        data_inds += list(data_unravel_indices[mask])
    proj_operator = sparse.coo_matrix((weights, (camera_inds, data_inds)))
    return proj_operator


def generate_synthetic_data(l):
    """Synthetic binary data"""
    rs = np.random.RandomState(0)
    n_pts = 36
    x, y = np.ogrid[0:l, 0:l]
    mask_outer = (x - l / 2.0) ** 2 + (y - l / 2.0) ** 2 < (l / 2.0) ** 2
    mask = np.zeros((l, l))
    points = l * rs.rand(2, n_pts)
    mask[(points[0]).astype(int), (points[1]).astype(int)] = 1
    mask = ndimage.gaussian_filter(mask, sigma=l / n_pts)
    res = np.logical_and(mask > mask.mean(), mask_outer)
    return np.logical_xor(res, ndimage.binary_erosion(res))

def Generate_synthetic_images_and_projections(l,alpha_l2,alpha_l1):

    # Generate synthetic images, and projections
    proj_operator = build_projection_operator(l, l // 7)
    data = generate_synthetic_data(l)
    proj = proj_operator @ data.ravel()[:, np.newaxis]
    proj += 0.15 * np.random.randn(*proj.shape)

    # Reconstruction with L2 (Ridge) penalization
    rgr_ridge = Ridge(alpha=alpha_l2)
    rgr_ridge.fit(proj_operator, proj.ravel())
    rec_l2 = rgr_ridge.coef_.reshape(l, l)

    # Reconstruction with L1 (Lasso) penalization
    # the best value of alpha was determined using cross validation
    # with LassoCV
    rgr_lasso = Lasso(alpha=alpha_l1)
    rgr_lasso.fit(proj_operator, proj.ravel())
    rec_l1 = rgr_lasso.coef_.reshape(l, l)
    
    fig = plt.figure(figsize=(8, 3.3))
    plt.subplot(131)
    plt.imshow(data, cmap=plt.cm.gray, interpolation="nearest")
    plt.axis("off")
    plt.title("original image")
    plt.subplot(132)
    plt.imshow(rec_l2, cmap=plt.cm.gray, interpolation="nearest")
    plt.title("L2 penalization")
    plt.axis("off")
    plt.subplot(133)
    plt.imshow(rec_l1, cmap=plt.cm.gray, interpolation="nearest")
    plt.title("L1 penalization")
    plt.axis("off")
    plt.subplots_adjust(hspace=0.01, wspace=0.01, top=1, bottom=0, left=0, right=1)

    fig.canvas.draw()
    image = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
    image = image.reshape(fig.canvas.get_width_height()[::-1] + (3,))
    plt.close(fig)
    return image

title = "Compressive sensing: Tomography reconstruction with L1 prior (Lasso)"

des="""This example illustrates the utilization of the Lasso algorithm to reconstruct an image from a collection of parallel projections taken at various angles. 
It also highlights the capability of the Lasso algorithm to accurately reconstruct images with no errors, even when noise is present in the projections.

The example demonstrates the effectiveness of the Lasso algorithm in reconstructing images using only a limited number of projections. 
In contrast, alternative reconstruction methods like the Ridge algorithm tend to introduce numerous labeling errors in the pixel reconstruction process"""

with gr.Blocks() as demo:
    gr.Markdown("# Compressive sensing: Tomography reconstruction with L1 prior (Lasso)")
    gr.Markdown("This demo is based on [Compressive sensing: Tomography reconstruction with L1 prior (Lasso)](https://scikit-learn.org/stable/auto_examples/applications/plot_tomography_l1_reconstruction.html#sphx-glr-auto-examples-applications-plot-tomography-l1-reconstruction-py).")
    gr.Markdown(f"{des}")
    
    with gr.Row():
        l=gr.Slider(minimum=100, maximum=500, step=1,  value = 128, label = "Linear size")
        alpha_l2=gr.Slider(minimum=0, maximum=1, step=0.001,  value = 0.2, label = "alpha l2")
        alpha_l1=gr.Slider(minimum=0, maximum=1, step=0.001,  value = 0.001, label = "alpha l1")



    output = gr.Image()
    btn = gr.Button(value="Submit")
    btn.click(Generate_synthetic_images_and_projections,[l,alpha_l2,alpha_l1],output)
    
demo.launch()