import gradio as gr
import matplotlib.pyplot as plt
import numpy as np

# Predefined hyperparameter sets
PARAM_SETS = {
    "Set A": {"param1": 0.1, "param2": 0.01, "param3": 100, "param4": 50},
    "Set B": {"param1": 0.2, "param2": 0.02, "param3": 200, "param4": 100}
}

def generate_plot(param1, param2, param3, param4):
    """Generate visualization based on hyperparameters"""
    plt.figure(figsize=(10, 6))
    x = np.linspace(0, 10, int(param3))
    y = np.sin(x * param1) * np.cos(x * param2) * param4
    plt.plot(x, y)
    plt.title(f'Parameter Visualization (p1={param1}, p2={param2}, p3={param3}, p4={param4})')
    plt.grid(True)
    return plt

def process_inputs(param_set, custom_param1, custom_param2, custom_param3, custom_param4, 
                  input1, input2):
    """Process inputs and return results"""
    # Determine which parameter set to use
    if param_set in PARAM_SETS:
        params = PARAM_SETS[param_set]
        p1, p2, p3, p4 = params.values()
    else:
        p1, p2, p3, p4 = custom_param1, custom_param2, custom_param3, custom_param4
    
    # Generate plot
    plot = generate_plot(p1, p2, p3, p4)
    
    # Calculate result (example calculation)
    result = (input1 * p1 + input2 * p2) * (p3 + p4)
    
    return plot, result

# Create interface
with gr.Blocks() as demo:
    gr.Markdown("# Hyperparameter Calculation and Visualization System")
    
    with gr.Row():
        with gr.Column():
            # Hyperparameter selection section
            param_set = gr.Dropdown(
                choices=["Custom"] + list(PARAM_SETS.keys()),
                value="Custom",
                label="Select Hyperparameter Set"
            )
            
            # Custom parameter inputs
            custom_param1 = gr.Number(value=0.1, label="Parameter 1 (Learning Rate)")
            custom_param2 = gr.Number(value=0.01, label="Parameter 2 (Weight Decay)")
            custom_param3 = gr.Number(value=100, label="Parameter 3 (Iterations)")
            custom_param4 = gr.Number(value=50, label="Parameter 4 (Batch Size)")
            
            # Input values
            input1 = gr.Number(value=1.0, label="Input Value 1")
            input2 = gr.Number(value=1.0, label="Input Value 2")
            
            submit_btn = gr.Button("Calculate")
        
        with gr.Column():
            # Output section
            plot_output = gr.Plot(label="Parameter Visualization")
            result_output = gr.Number(label="Calculation Result")
    
    # Auto-fill parameters when selecting predefined sets
    def update_params(param_set):
        if param_set in PARAM_SETS:
            params = PARAM_SETS[param_set]
            return [params["param1"], params["param2"], params["param3"], params["param4"]]
        return [gr.skip(), gr.skip(), gr.skip(), gr.skip()]
    
    param_set.change(
        update_params,
        inputs=[param_set],
        outputs=[custom_param1, custom_param2, custom_param3, custom_param4]
    )
    
    # Submit button event
    submit_btn.click(
        process_inputs,
        inputs=[param_set, custom_param1, custom_param2, custom_param3, custom_param4,
                input1, input2],
        outputs=[plot_output, result_output]
    )

# Launch application
demo.launch()