import gradio as gr
import numpy as np
from pylops.signalprocessing import Radon2D
from typing import Literal

def compute_radon_fn(image,
                     kind: Literal["linear", "hyperbolic", "parabolic"] = "linear"):
    # Convert image to grayscale if needed.
    if image.ndim == 3:
        image_gray = np.mean(image, axis=2)
    else:
        image_gray = image

    nt, nh = image_gray.shape
    # Define axes centered around zero.
    taxis = np.linspace(-nt/2, nt/2, nt)   # time axis (rows)
    haxis = np.linspace(-nh/2, nh/2, nh)   # spatial axis (columns)
    # Set detector axis (pxaxis) to cover the full image diagonal.
    # npx = int(np.ceil(np.sqrt(nt**2 + nh**2)))
    npx = nh # why?
    pxaxis = np.linspace(-npx/2, npx/2, npx)

    print(f"Shapes:\n taxis:{taxis.shape}, haxis:{haxis.shape}, pxaxis:{pxaxis.shape}")

    # Create the Radon2D operator with engine 'numba' and centeredh=True.
    R = Radon2D(taxis, haxis, pxaxis,
                kind=kind, centeredh=True, interp=True,
                onthefly=False, engine='numpy', dtype='float64', name='R')

    # Compute the forward radon transform.
    radon_data_flat = R.dot(image_gray.flatten())
    # Deduce the number of projections from the operator shape.
    nproj = R.shape[0] // npx
    # Reshape to (nproj, L) so each row corresponds to one projection.
    radon_data = radon_data_flat.reshape((nproj, npx))
    
    # Transpose for display: p (detector coordinate) on x-axis, τ on y-axis.
    radon_display = radon_data.T
    # Normalize for display.
    radon_display = (radon_display - radon_display.min()) / (radon_display.max() - radon_display.min() + 1e-8)
    
    # Save the state including the radon data and operator for inverse computation.
    state = {
         "radon_data": radon_data,     # shape: (nproj, L)
         "image_shape": image_gray.shape,
         "R": R,                       # the Radon2D operator
         "L": npx                      # detector length
    }
    return radon_display, state

def apply_mask_and_inverse_fn(state, mask_image):
    if mask_image is None:
        return None

    # Ensure mask is single-channel.
    if mask_image.ndim == 3:
        mask_image = mask_image[..., 0]

    radon_data = state["radon_data"]
    image_shape = tuple(state["image_shape"])
    L = state["L"]
    
    # The displayed radon image was transposed, so transpose mask back.
    mask = mask_image.T
    # Create binary mask: painted pixels (value > 0.5) become 1.
    mask_binary = (mask > 0.5).astype(float)
    
    # Apply the mask: zero-out masked pixels in the radon data.
    radon_masked = radon_data * (1 - mask_binary)
    
    # Reconstruct the image using the adjoint (transpose) of the operator.
    R = state["R"]
    rec_flat = R.T.dot(radon_masked.flatten())
    rec = rec_flat.reshape(image_shape)
    # Normalize reconstruction for display.
    rec_norm = (rec - rec.min()) / (rec.max() - rec.min() + 1e-8)
    return rec_norm

with gr.Blocks() as demo:
    gr.Markdown("## Radon Transform with Interactive Masking (Using PyLops Radon2D)")
    
    with gr.Row():
        image_input = gr.Image(label="Input Image", type="numpy")
    
    compute_button = gr.Button("Compute Radon")
    radon_output = gr.Image(label="Radon Transform Image", interactive=False, type="numpy")
    # Use the updated ImageEditor component for interactive masking.
    radon_drawing = gr.ImageEditor(label="Paint on Radon (mask out pixels)", type="numpy")
    inverse_output = gr.Image(label="Reconstructed Image", interactive=False, type="numpy")
    
    # State to hold radon data and operator.
    state = gr.State()
    
    # When "Compute Radon" is clicked, compute the radon transform.
    compute_button.click(
        fn=compute_radon_fn,
        inputs=[image_input],
        outputs=[radon_output, state]
    )
    
    # When the user edits (paints) the radon image, apply the mask and compute the inverse.
    radon_drawing.change(
        fn=apply_mask_and_inverse_fn,
        inputs=[state, radon_drawing],
        outputs=[inverse_output]
    )
    
    gr.Markdown(
        "**Instructions:** Upload an image and click 'Compute Radon' to compute the radon transform using PyLops’ Radon2D (with engine 'numba' and centeredh=True). Then use the paintbrush tool to mask parts of the radon image and see the resulting reconstruction."
    )

demo.launch()