import gradio as gr
import json
import pandas as pd
from Engine import Engine

def run_study(mode, benchmark_func, optimizers, dim, dataset, epochs, batch_size, lr, use_sa, sa_temp, sa_cooling_rate):
    # Ensure optimizers is a list
    optimizers = [optimizers] if isinstance(optimizers, str) else optimizers or []
    
    if not optimizers:
        raise gr.Error("Please select at least one optimizer.")
    if mode == "Benchmark Optimization" and not benchmark_func:
        raise gr.Error("Please select a benchmark function.")
    if mode == "ML Task Training" and not dataset:
        raise gr.Error("Please select a dataset.")

    config = {
        'mode': 'benchmark' if mode == "Benchmark Optimization" else 'ml_task',
        'benchmark_func': benchmark_func,
        'optimizers': optimizers,
        'dim': int(dim) if dim else 2,
        'dataset': dataset,
        'epochs': int(epochs) if epochs else 10,
        'batch_size': int(batch_size) if batch_size else 32,
        'lr': float(lr) if lr else 0.001,
        'use_sa': use_sa if 'AzureSky' in optimizers else None,
        'sa_temp': float(sa_temp) if 'AzureSky' in optimizers and use_sa else None,
        'sa_cooling_rate': float(sa_cooling_rate) if 'AzureSky' in optimizers and use_sa else None,
        'max_iter': 100
    }
    runner = Engine()
    results = runner.run(config)

    if config['mode'] == 'benchmark':
        metrics_df = pd.DataFrame(results['metrics'], index=config['optimizers'])
        return results['plot'], None, metrics_df, json.dumps(results, indent=2), "Study completed successfully."
    else:
        metrics_df = pd.DataFrame(results['metrics'], index=config['optimizers'])
        return results['plot_acc'], results['plot_loss'], metrics_df, json.dumps(results, indent=2), "Study completed successfully."

def export_results(results_json):
    with open("results.json", "w") as f:
        f.write(results_json)
    return "results.json"

def toggle_azure_settings(optimizers):
    optimizers = [optimizers] if isinstance(optimizers, str) else optimizers or []
    return gr.update(visible='AzureSky' in optimizers)

def toggle_tabs(mode):
    return gr.update(visible=mode == 'Benchmark Optimization'), gr.update(visible=mode == 'ML Task Training')

with gr.Blocks(theme=gr.themes.Soft(), title="Nexa R&D Studio", css="""
    .gr-button { margin-top: 10px; }
    .gr-box { border-radius: 8px; }
    .status-message { color: green; font-weight: bold; }
""") as app:
    gr.Markdown("""
        # Nexa R&D Studio
        A visual research tool for comparing and evaluating optimizers on benchmark functions and ML tasks.
        Select a mode, configure your study, and analyze results with interactive plots and metrics.
    """)

    with gr.Tabs():
        with gr.TabItem("Study Configuration"):
            mode = gr.Radio(
                ['Benchmark Optimization', 'ML Task Training'],
                label='Study Mode',
                value='Benchmark Optimization',
                info='Choose between optimizing benchmark functions or training on ML datasets.'
            )

            with gr.Row():
                with gr.Column():
                    optimizers = gr.CheckboxGroup(
                        ['AzureSky', 'Adam', 'SGD', 'AdamW', 'RMSprop'],
                        label='Optimizers',
                        info='Select optimizers to compare. AzureSky includes a Simulated Annealing option.'
                    )
                    with gr.Accordion("AzureSky Ablation Settings", open=False, visible=False) as azure_settings:
                        use_sa = gr.Checkbox(
                            label='Enable Simulated Annealing (AzureSky)',
                            value=True,
                            info='Toggle Simulated Annealing for AzureSky optimizer.'
                        )
                        sa_temp = gr.Number(
                            label='Initial SA Temperature',
                            value=1.0,
                            minimum=0.1,
                            info='Controls exploration in Simulated Annealing (higher = more exploration).'
                        )
                        sa_cooling_rate = gr.Number(
                            label='SA Cooling Rate',
                            value=0.95,
                            minimum=0.1,
                            maximum=0.99,
                            info='Rate at which SA temperature decreases (closer to 1 = slower cooling).'
                        )

                with gr.Column():
                    with gr.Group(visible=True) as benchmark_tab:
                        benchmark_func = gr.Dropdown(
                            ['Himmelblau', 'Ackley', 'Adjiman', 'Brent'],
                            label='Benchmark Function',
                            info='Select a mathematical function to optimize.'
                        )
                        dim = gr.Number(
                            label='Dimensionality',
                            value=2,
                            minimum=2,
                            info='Number of dimensions for the benchmark function.'
                        )
                    with gr.Group(visible=False) as ml_task_tab:
                        dataset = gr.Dropdown(
                            ['MNIST', 'CIFAR-10'],
                            label='Dataset',
                            info='Select a dataset for ML training.'
                        )
                        epochs = gr.Number(
                            label='Epochs',
                            value=10,
                            minimum=1,
                            info='Number of training epochs.'
                        )
                        batch_size = gr.Number(
                            label='Batch Size',
                            value=32,
                            minimum=1,
                            info='Number of samples per training batch.'
                        )
                        lr = gr.Number(
                            label='Learning Rate',
                            value=0.001,
                            minimum=0,
                            info='Learning rate for optimizers.'
                        )

            run_button = gr.Button('Run Study', variant='primary')

        with gr.TabItem("Results"):
            status_message = gr.Markdown("Configure and run a study to view results.", elem_classes=["status-message"])
            with gr.Row():
                plot1 = gr.Plot(label='Main Plot (Benchmark or Accuracy)')
                plot2 = gr.Plot(label='Loss Plot (ML Mode)')
            metrics_df = gr.Dataframe(label='Metrics Table')
            metrics_json = gr.JSON(label='Detailed Metrics')
            export_button = gr.Button('Export Results as JSON')
            export_file = gr.File(label='Download Results')

    mode.change(
        fn=toggle_tabs,
        inputs=mode,
        outputs=[benchmark_tab, ml_task_tab]
    )
    optimizers.change(
        fn=toggle_azure_settings,
        inputs=optimizers,
        outputs=azure_settings
    )
    run_button.click(
        fn=run_study,
        inputs=[mode, benchmark_func, optimizers, dim, dataset, epochs, batch_size, lr, use_sa, sa_temp, sa_cooling_rate],
        outputs=[plot1, plot2, metrics_df, metrics_json, status_message]
    )
    export_button.click(
        fn=export_results,
        inputs=metrics_json,
        outputs=export_file
    )

# Launch without share parameter for Hugging Face Spaces
app.launch()