import gradio as gr
import torch
import numpy as np

def apply_desaturation(image, factor, method):
    if image is None:
        return None
        
    # Convert to torch tensor and normalize
    image = torch.from_numpy(image).float() / 255.0
    
    if method == "luminance (Rec.709)":
        grayscale = 0.2126 * image[..., 0] + 0.7152 * image[..., 1] + 0.0722 * image[..., 2]
    elif method == "luminance (Rec.601)":
        grayscale = 0.299 * image[..., 0] + 0.587 * image[..., 1] + 0.114 * image[..., 2]
    elif method == "average":
        grayscale = image.mean(dim=2)
    elif method == "lightness":
        grayscale = (torch.max(image, dim=2)[0] + torch.min(image, dim=2)[0]) / 2
    
    # Blend original and grayscale
    grayscale = (1.0 - factor) * image + factor * grayscale.unsqueeze(-1).repeat(1, 1, 3)
    
    # Convert back to uint8
    output = (grayscale.clamp(0, 1).numpy() * 255).astype(np.uint8)
    
    return output

def apply_posterize(image, threshold):
    if image is None:
        return None
        
    # Convert to torch tensor and normalize
    image = torch.from_numpy(image).float() / 255.0
    
    # Convert to grayscale and apply threshold
    grayscale = image.mean(dim=2, keepdim=True)
    posterized = (grayscale > threshold).float()
    
    # Repeat for all channels
    posterized = posterized.repeat(1, 1, 3)
    
    # Convert back to uint8
    output = (posterized.numpy() * 255).astype(np.uint8)
    
    return output

def create_effects_tab():
    with gr.Tab("Simple Effects"):
        
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="Input Image", height=256)
                
                with gr.Tabs():
                    with gr.Tab("Desaturate"):
                        desaturate_method = gr.Dropdown(
                            choices=[
                                "luminance (Rec.709)",
                                "luminance (Rec.601)",
                                "average",
                                "lightness"
                            ],
                            value="luminance (Rec.709)",
                            label="Method"
                        )
                        desaturate_factor = gr.Slider(
                            minimum=0.0,
                            maximum=1.0,
                            value=1.0,
                            step=0.05,
                            label="Factor"
                        )
                        desaturate_btn = gr.Button("Apply Desaturation")
                        
                    with gr.Tab("Posterize"):
                        threshold = gr.Slider(
                            minimum=0.0,
                            maximum=1.0,
                            value=0.5,
                            step=0.05,
                            label="Threshold"
                        )
                        posterize_btn = gr.Button("Apply Posterization")
            
            with gr.Column():
                output_image = gr.Image(label="Processed Image")
        
        desaturate_btn.click(
            fn=apply_desaturation,
            inputs=[input_image, desaturate_factor, desaturate_method],
            outputs=output_image
        )
        
        posterize_btn.click(
            fn=apply_posterize,
            inputs=[input_image, threshold],
            outputs=output_image
        )