import pandas as pd

from pathlib import Path
import typing as tp
import gradio as gr

audio_attacks_with_variations = ['random_noise', 'lowpass_filter', 'highpass_filter', 'boost_audio', 'duck_audio', 'shush_fraction']
audio_models = ['wavmark', 'timbre', 'audioseal']
audio_metrics = ['snr', 'sisnr', 'stoi', 'ba', 'pesq', 'detect_prob']

image_attacks_with_variations = [
    # "crop",
    "jpeg",
    "brightness",
    "contrast",
    "saturation",
    "sharpness",
    "resize",
    "perspective",
    "median_filter",
    "hue",
    "gaussian_blur",
]
image_models = ["dctdwt", "fnns", "hidden", "ssl", "trustmark", "wam"]
image_metrics = [
    "psnr",
    "ssim",
    "lpips",
    "bit_acc",
    "p_value",
    "word_acc",
    "watermark_det_score",
]
    
def plot_data(metric, selected_attack, all_attacks_df):  
    attack_df = all_attacks_df[all_attacks_df.attack == selected_attack]
    
    if metric == "None":
        return gr.LinePlot(x_bin=None)
        
    return gr.LinePlot(
            attack_df,
            x="strength",
            y=metric,
            color="model",
        )

def mk_audio_variations(csv_file: Path, ):
    all_attacks_df = pd.read_csv(csv_file)
    
    with gr.Row():
        group_by = gr.Radio(audio_metrics, value=audio_metrics[0], label="Choose metric")
        attacks_dropdown = gr.Dropdown(
            audio_attacks_with_variations, label=audio_attacks_with_variations[0], info="Select attack"
        )
    
    attacks_by_strength = plot_data(group_by.value, attacks_dropdown.value, all_attacks_df)

    all_graphs = [attacks_by_strength, ]
    
    group_by.change(
        lambda group: plot_data(group, attacks_dropdown.value, all_attacks_df),
        group_by,
        all_graphs
    )
    
    attacks_dropdown.change(
        lambda attack: plot_data(group_by.value, attack, all_attacks_df),
        attacks_dropdown,
        all_graphs
    )


def mk_image_variations(csv_file: Path, ):
    all_attacks_df = pd.read_csv(csv_file)
    
    with gr.Row():
        group_by = gr.Radio(image_metrics, value=image_metrics[0], label="Choose metric")
        attacks_dropdown = gr.Dropdown(
            image_attacks_with_variations, label=image_attacks_with_variations[0], info="Select attack"
        )
    
    attacks_by_strength = plot_data(group_by.value, attacks_dropdown.value, all_attacks_df)

    all_graphs = [attacks_by_strength, ]
    
    group_by.change(
        lambda group: plot_data(group, attacks_dropdown.value, all_attacks_df),
        group_by,
        all_graphs
    )
    
    attacks_dropdown.change(
        lambda attack: plot_data(group_by.value, attack, all_attacks_df),
        attacks_dropdown,
        all_graphs
    )