import gradio as gr
import plotly.graph_objects as go
from huggingface_hub import hf_hub_download
from datasets import load_dataset
import pathlib
import json
import pandas as pd
from evaluation import load_boundary, load_boundaries
from constellaration import forward_model, initial_guess
from constellaration.boozer import boozer
from constellaration.utils import (
    file_exporter,
    visualization,
    visualization_utils,
)


organization = 'proxima-fusion'
results_repo = f'{organization}/constellaration-bench-results'

def read_result_from_hub(filename):
    local_path = hf_hub_download(
        repo_id=results_repo,
        repo_type="dataset",
        filename=filename,
    )
    return local_path

def make_boundary_plot(boundary):
    boundary_ax = visualization.plot_boundary(boundary)
    boundary_plot = boundary_ax.get_figure()
    return boundary_plot

def make_interactive_plot(boundary):
    interactive_plot = visualization.plot_surface(boundary)
    return interactive_plot

def make_boozer_plot(boundary):
    settings = forward_model.ConstellarationSettings.default_high_fidelity_skip_qi()
    boundary_metrics, boundary_equilibrium = forward_model.forward_model(
        boundary, settings=settings
    )
    boozer_settings = boozer.BoozerSettings(normalized_toroidal_flux=[0.0, 0.25, 1])
    boozer_plot = visualization.plot_boozer_surfaces(
        boundary_equilibrium, settings=boozer_settings
    )
    return

def make_flux_surface_plot(boundary):
    settings = forward_model.ConstellarationSettings.default_high_fidelity_skip_qi()
    boundary_metrics, boundary_equilibrium = forward_model.forward_model(
        boundary, settings=settings
    )
    flux_surface_plot = visualization.plot_flux_surfaces(
        boundary_equilibrium, boundary
    )
    return flux_surface_plot

def gradio_interface() -> gr.Blocks:
    with gr.Blocks() as demo:
        gr.Markdown("""
                    # Welcome to the ConStellaration Boundary Explorer!

                    ### Here, you can visualize submissions to the ConStellaration Leaderboard, generate and visualize new random boundaries, or upload and visualize your own!

                    If you want to explore additional visualizations, like Boozer plots and flux surface plots, check out the `constellaration` library's [visualization tools](https://github.com/proximafusion/constellaration/blob/main/notebooks/boundary_explorer.ipynb).
                    """)
        
        ds = load_dataset(results_repo, split='train')
        full_df = pd.DataFrame(ds)
        filenames = full_df['result_filename'].to_list()

        boundary = gr.State()

        mode_selector = gr.Radio(choices=["Generate", "Leaderboard", "Upload"],
                                 label="Select input method:",
                                 value="Leaderboard")


        with gr.Row():
            with gr.Column(visible=False) as generate_ui:
                aspect_ratio = gr.Number(label="Aspect Ratio", value=3)
                elongation = gr.Number(label="Elongation", value=0.5)
                rotational_transform = gr.Number(label="Rotational Transform", value=0.4)
                n_field_periods = gr.Number(label="Number of Period Fields", value=3)
                generate_btn = gr.Button(value="Generate")

            with gr.Column(visible=True) as leaderboard_ui:
                dropdown = gr.Dropdown(choices=filenames, label="Choose a leaderboard entry", value=filenames[0])
                rld_btn = gr.Button(value="Reload")

            with gr.Column(visible=False) as upload_ui:
                upload_box = gr.File(file_types=[".json"], label="Upload your boundary file")
            
            
            with gr.Column(visible=True):
                interactive_plot = gr.Plot()
                boundary_plot = gr.Plot()
            #    boozer_plot = gr.Plot()
            #    flux_surface_plot = gr.Plot()

        def update_ui(mode):
            return (
                gr.update(visible=(mode == "Generate")),
                gr.update(visible=(mode == "Leaderboard")),
                gr.update(visible=(mode == "Upload")),  
            )

        mode_selector.change(update_ui, inputs=[mode_selector], outputs=[generate_ui, leaderboard_ui, upload_ui])

        def get_boundary_from_leaderboard(selected_file):
            row = full_df[full_df['result_filename'] == selected_file].iloc[0]
            if row['problem_type'] == 'mhd_stable':
                raise gr.Error("Sorry this isn't implemented for files with multiple boundaries yet!")
            else:
                boundary = load_boundary(row['boundary_json'])

            return boundary

        # dropdown.change(get_boundary_from_leaderboard, dropdown, [boundary_plot, interactive_plot, boozer_plot, flux_surface_plot])
        rld_btn.click(get_boundary_from_leaderboard, dropdown, boundary) \
            .then(make_interactive_plot, boundary, interactive_plot) \
            .then(make_boundary_plot, boundary, boundary_plot) \
            # .then(make_boozer_plot, boundary, boozer_plot) \
            # .then(make_flux_surface_plot, boundary, flux_surface_plot)


        def get_boundary_vis_from_upload(uploaded_file):
            if uploaded_file is None:
                raise gr.Error("Please upload a file.")
            with open(uploaded_file.name, 'r') as f:
                data = f.read()
            boundary = load_boundary(data)
            return boundary
        
        upload_box.change(get_boundary_vis_from_upload, inputs=[upload_box], outputs=boundary) \
            .then(make_interactive_plot, boundary, interactive_plot) \
            .then(make_boundary_plot, boundary, boundary_plot) \
            # .then(make_boozer_plot, boundary, boozer_plot) \
            # .then(make_flux_surface_plot, boundary, flux_surface_plot)

        def generate_random_boundary(aspect_ratio, elongation, rotational_transform, n_field_periods):
            boundary = initial_guess.generate_rotating_ellipse(
                aspect_ratio=aspect_ratio, elongation=elongation, rotational_transform=rotational_transform, n_field_periods=n_field_periods
            )
            return boundary
        
        generate_btn.click(generate_random_boundary, [aspect_ratio, elongation, rotational_transform, n_field_periods], boundary) \
            .then(make_interactive_plot, boundary, interactive_plot) \
            .then(make_boundary_plot, boundary, boundary_plot) \
            # .then(make_boozer_plot, boundary, boozer_plot) \
            # .then(make_flux_surface_plot, boundary, flux_surface_plot)

    return demo


if __name__ == "__main__":
    gradio_interface().launch()