import gradio as gr
from simulations import finance_demo, quantum_demo, fluid_demo, bio_demo
from poseidon_model import load_model, run_inference_by_domain, run_inference_on_dataset, plot_output, plot_comparison


def run_poseidon_demo(domain, contrast, cmap):
    """
        Loads the POSEIDON model and runs it on synthetic input data
        based on the selected scientific domain (e.g., Finance, Quantum).

        Args:
            domain (str): Selected scientific field.
            contrast (float): Contrast setting for visualization.
            cmap (str): Colormap choice for heatmap.

        Returns:
            Matplotlib figure showing the output.
    """
    model = load_model()
    output = run_inference_by_domain(model, domain)
    return plot_output(output, contrast=contrast, cmap=cmap)



def render_demo(domain):
    """
        Returns a mini-simulation plot and a descriptive explanation
        for the selected domain.

        Args:
            domain (str): One of Finance, Quantum, Fluid Dynamics, Biology.

        Returns:
            Tuple of (plot, explanatory markdown string).
    """

    if domain == "Finance":
        return finance_demo(), (
            "📍 **Finance:** PDEs like Black-Scholes are used to model option pricing. "
            "Imagine fine-tuning Poseidon to forecast derivatives across market regimes!"
        )
    elif domain == "Quantum":
        return quantum_demo(), (
            "📍 **Quantum Mechanics:** Schrödinger's equation is a core PDE in quantum physics. "
            "Could Poseidon learn to generalize across quantum systems?"
        )
    elif domain == "Fluid Dynamics":
        return fluid_demo(), (
            "📍 **Fluid Dynamics:** Poseidon is pretrained here! This sim shows 1D flow, "
            "but Poseidon can do much more."
        )
    elif domain == "Biology / Medicine":
        return bio_demo(), (
            "📍 **Biology:** Reaction-diffusion equations appear in tissue growth and morphogenesis. "
            "Poseidon could help model organ behavior!"
        )
    else:
        return None, "Pick a domain to explore how Poseidon might apply!"



def run_poseidon_real_dataset(dataset_name):
    """
        Loads Poseidon and runs inference on a real scientific dataset from the Hub.

        Args:
            dataset_name (str): Dataset ID from dropdown.

        Returns:
            Matplotlib figure with side-by-side comparison of input vs output.
    """
    model = load_model()
    input_array, output_array = run_inference_on_dataset(model, dataset_name)
    return plot_comparison(input_array, output_array)


# this part defines the app
with gr.Blocks() as demo:
    gr.Markdown("# 🔱 POSEIDON Application Across Scientific Domains 🔱")

    gr.Markdown("### **Welcome to the POSEIDON Playground!**")
    gr.Markdown("Ever dreamed of solving physics equations with a single click? You’re in the right place.")
    gr.Markdown("POSEIDON is a foundation model that learned to solve partial differential equations (PDEs) — the magical "
                "math behind fluid flows, quantum mechanics, financial markets, and even biology!")

    gr.Markdown("## ☑️ 1. Pick a scientific domain to see a simple PDE simulation and Explanation.")
    domain_dropdown = gr.Dropdown(
        ["Finance", "Quantum", "Fluid Dynamics", "Biology / Medicine"],
        label="Choose The Field",
        value="Finance"
    )
    sim_output = gr.Plot()
    sim_text = gr.Markdown()

    domain_dropdown.change(fn=render_demo, inputs=domain_dropdown, outputs=[sim_output, sim_text])

    gr.Markdown("## 🚀 Run a test output from the POSEIDON model based on the chosen domain")

    with gr.Row():
        gr.Markdown("Play with contrast and choose the colormap you prefer.")
        contrast_slider = gr.Slider(0.5, 5.0, value=2.0, step=0.1, label=" Contrast")
        cmap_dropdown = gr.Dropdown(
            ["inferno", "viridis", "plasma"],
            label="Colormap",
            value="inferno"
        )

    with gr.Row():
        poseidon_button = gr.Button("POSEIDON Test Output")
        

    poseidon_plot = gr.Plot()
    poseidon_button.click(
        fn=run_poseidon_demo,
        inputs=[domain_dropdown, contrast_slider, cmap_dropdown],
        outputs=poseidon_plot
    )

    gr.Markdown("# ")

    gr.Markdown("## ☑️ 2. Try POSEIDON on Real Scientific Datasets")
    dataset_dropdown = gr.Dropdown(
        ["fluids.incompressible.Sines", "fluids.compressible.Riemann", "reaction_diffusion.AllenCahn"],
        label="Choose a Real Dataset"
    )
    dataset_button = gr.Button("POSEIDON on Dataset")
    dataset_plot = gr.Plot()

    dataset_button.click(
        fn=run_poseidon_real_dataset,
        inputs=[dataset_dropdown],
        outputs=dataset_plot
    )



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