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 shows how to reconstruct an image from a set of parallel projections, acquired along different angles, using the Lasso algorithm. 
The Lasso algorithm is a type of compressed sensing algorithm that uses prior information about the sparsity of the image to reconstruct it from a small number of projections. The example shows that the Lasso algorithm can successfully reconstruct images with zero error, even if noise was added to the projections. 
In comparison, other methods of reconstruction, such as the Ridge algorithm, produce a large number of labeling errors for the pixels.
The example demonstrates the effectiveness of the Lasso algorithm for image reconstruction from a small number of projections. This is a promising technique for applications where it is not possible or practical to acquire a large number of projections, such as in medical imaging or in astronomy."""

with gr.Blocks(title=title) as demo:
    gr.Markdown(f"#{title}")
    gr.Markdown(f"{des}")
    gr.Markdown("This demo is based on this [scikit-learn example](https://scikit-learn.org/stable/auto_examples/applications/plot_tomography_l1_reconstruction.html#sphx-glr-auto-examples-applications-plot-tomography-l1-reconstruction-py).")
    
    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(fn=Generate_synthetic_images_and_projections, inputs = [l,alpha_l2,alpha_l1], outputs=output)

demo.launch()