import gradio as gr
print("Gradio version:", gr.__version__)
import os
import matplotlib.pyplot as plt
import netCDF4 as nc
from huggingface_hub import hf_hub_download
from utils import fig2img, create_gif_from_frames

# Global dictionary to store datasets in memory
datasets = {}
dataset_metadata = {"CE-RPUI": {"key":"data", "features":["density", "horz_velocity", "vertical_velocity", "pressure", "energy"]}, "CE-RM": {"key":"solution", "features":["density", "horz_velocity", "vertical_velocity", "pressure", "energy"]}, "NS-PwC": {"key":"velocity", "features":["horz_velocity", "vertical_velocity", "passive_tracer"]}}
feature_label_map = {"density": "Density", "horz_velocity": "Horizontal Velocity", "vertical_velocity": "Vertical Velocity", "pressure": "Pressure", "energy": "Energy", "passive_tracer": "Passive Tracer"}

def load_dataset_into_memory(dataset_name):
    if dataset_name not in datasets:
        print(f"Loading dataset: {dataset_name}...")

        dataset_path = hf_hub_download(
            repo_id=f"camlab-ethz/{dataset_name}",
            filename=dataset_metadata[dataset_name]['key'] + "_0.nc",
            repo_type="dataset",
            token=os.getenv("HF_TOKEN")
        )

        dataset = nc.Dataset(dataset_path)
        datasets[dataset_name] = dataset.variables[dataset_metadata[dataset_name]['key']]  # Store only the data variable

    return datasets[dataset_name]

# Function to generate GIFs dynamically
def generate_gifs(dataset_name, index):
    data = load_dataset_into_memory(dataset_name)  # Load dataset from memory
    first_traj = data[index]  # Extract trajectory at given index

    # Extract features
    if dataset_name == "NS-PwC":
        features = {
            "horz_velocity": first_traj[:, 0, :, :],
            "vertical_velocity": first_traj[:, 1, :, :],
            "passive_tracer": first_traj[:, 2, :, :],
        }
    else:
        features = {
            "density": first_traj[:, 0, :, :],
            "horz_velocity": first_traj[:, 1, :, :],
            "vertical_velocity": first_traj[:, 2, :, :],
            "pressure": first_traj[:, 3, :, :],
            "energy": first_traj[:, 4, :, :]
        }

    gif_paths = []
    output_dir = "episode_gifs"
    os.makedirs(output_dir, exist_ok=True)  # Ensure output directory exists

    for feature_name, feature in features.items():
        gif_filename = f"{output_dir}/{dataset_name.lower()}_{feature_name}_{index}.gif"

        # Only generate if the GIF doesn't already exist
        if not os.path.exists(gif_filename):
            print(f"Generating GIF: {gif_filename}")
            frames = []
            for i in range(len(feature)):
                fig = plt.figure(figsize=(1.28, 1.28), dpi=100, frameon=False)
                ax = plt.Axes(fig, [0., 0., 1., 1.])
                ax.set_axis_off()
                fig.add_axes(ax)
                
                ax.imshow(feature[i, :, :])
                frames.append(fig2img(fig))
                plt.close(fig)

            create_gif_from_frames(frames, gif_filename)
        else:
            print(f"Using cached GIF: {gif_filename}")

        gif_paths.append(gif_filename)

    return gif_paths  # Return paths to update the Gradio UI


gif_css = """
#gif-label {
    text-align: center;
    width: 100%;
    display: flex;
    justify-content: center;
}
"""

css = """
#label {
    text-align: center;
    width: 100%;
}
"""

default_dataset = "CE-RPUI"
default_index = 0
default_gif_paths = generate_gifs(default_dataset, default_index)
default_features = dataset_metadata[default_dataset]["features"]

with gr.Blocks(css=css) as demo:
    gr.Markdown(
        """
        # POSEIDON: Foundation Models for PDEs 🌊🔬

        POSEIDON is a foundation model for solving Partial Differential Equations (PDEs) efficiently. Instead of training a separate model for each PDE, POSEIDON learns a general solution operator—allowing it to generalize across different physics with minimal data. Think of it as the GPT4.5 for PDEs, trained on a diverse set of fluid dynamics equations and capable of adapting to new, unseen physical systems.

        # **Dataset Explorer**
        POSEIDON provides solutions to a variety of fluid dynamics problems. Below are a few datasets you can explore:

        ### **CE-RM (Richtmyer-Meshkov)**
        - Based on the compressible Euler equations, this dataset models shock-driven instability at fluid interfaces.
        - Used in astrophysics, fusion research, and high-speed aerodynamics.

        ### **NS-PwC (Navier-Stokes - Piecewise Constant Vorticity)**
        - Based on the incompressible Navier-Stokes equations, modeling turbulence and vortex dynamics.
        - Applications include climate modeling, aerodynamics, and turbulent flow control.

        ### **CE-RPUI (Riemann Problems with Uncertain Interfaces)**
        - Models shock interactions across uncertain boundaries, crucial for hypersonic flows and uncertainty quantification.
        - Helps in high-speed aerodynamics and robust PDE solvers.

        Explore these datasets to visualize fluid behavior and analyze dynamic flow evolution!
        """
    )

    with gr.Row():
        dataset_selector = gr.Dropdown(choices=["CE-RPUI", "CE-RM", "NS-PwC"], value="CE-RPUI", label="Select Dataset")
        index_slider = gr.Slider(minimum=0, maximum=100, value=0, step=1, label="Select Trajectory Index")

    gif_outputs = []
    label_outputs = []
    with gr.Row(equal_height=True) as gif_container:
        for i in range(5):
            with gr.Column(scale=1, min_width=0):
                if i < len(default_features):
                    output = gr.Image(value=default_gif_paths[i], type="filepath", show_label=False, container=False)
                    label = gr.Markdown(f"**{feature_label_map[default_features[i]]}**", elem_id="label")
                else:
                    output = gr.Image(visible=False, type="filepath", show_label=False, container=False)
                    label = gr.Markdown(visible=False, elem_id="label")

                gif_outputs.append(output)
                label_outputs.append(label)


    def update_layout(dataset_name, index):
        features = dataset_metadata[dataset_name]["features"]
        gif_paths = generate_gifs(dataset_name, index)

        output_values = []
        label_values = []
        for i in range(5):
            if i < len(features):
                output_values.append(gr.update(value=gif_paths[i], visible=True))
                label_values.append(gr.update(value=f"**{feature_label_map[features[i]]}**", visible=True))
            else:
                output_values.append(gr.update(visible=False))
                label_values.append(gr.update(visible=False))

        return output_values + label_values

    dataset_selector.change(update_layout, inputs=[dataset_selector, index_slider], outputs=gif_outputs + label_outputs)
    index_slider.change(update_layout, inputs=[dataset_selector, index_slider], outputs=gif_outputs + label_outputs)

    gr.Markdown(
        """
        ## Key Innovations
        ### **Multiscale Operator Transformer (scOT)**  
        A hierarchical transformer-based architecture that captures PDE solution dynamics across multiple spatial and temporal scales. It uses shifted-window attention (SwinV2) to efficiently process large solution spaces.

        ### **Continuous-in-Time Learning**  
        Instead of learning PDE solutions at discrete time steps, POSEIDON uses time-conditioned layer normalization, enabling predictions at **any arbitrary time**—like a **true continuous function**.

        ### **All2All Training Strategy**  
        By leveraging the semi-group property of PDEs, POSEIDON scales training data quadratically without additional simulations. Every time step becomes a learning opportunity!

        ### **Pretrained on Fluid Dynamics, Generalizes to New Physics**  
        Trained on compressible Euler and Navier-Stokes equations, POSEIDON transfers its knowledge to unseen wave, diffusion, and reaction-diffusion PDEs—a huge step for scientific machine learning!

        ### **Outperforms FNO & Neural Operators**  
        POSEIDON achieves the same accuracy as an FNO trained on 1024 samples—using only 20 samples. That's a 50x efficiency boost in sample efficiency.

        ---

        ## Why Does This Matter?  
        Traditional PDE solvers are computationally expensive. POSEIDON is a general-purpose neural PDE solver that:  

        • Works across multiple physics domains  
        • Requires fewer training samples  
        • Enables real-time simulation & forecasting  

        It's a step towards universal scientific models, just like foundation models transformed NLP and vision.

        ---

        ## Try POSEIDON Now!  

        You can experiment and empower your research with POSEIDON-T (21M parameters), POSEIDON-B (158M parameters), and POSEIDON-L (629M parameters).

        • **Pretrained models & datasets**: [Hugging Face Hub](https://huggingface.co/camlab-ethz)  
        • **Code & Paper**: [GitHub](https://github.com/camlab-ethz/poseidon) | [arXiv](https://arxiv.org/abs/2405.19101)  
        • **Join the Discussion**: [Hugging Face Forums](https://discuss.huggingface.co/)  

        Let's reshape the future of PDE solving—one foundation model at a time!

        ---

        If you use POSEIDON in your research, please cite the paper:
        ```
        @misc{herde2024poseidon,
            title={Poseidon: Efficient Foundation Models for PDEs}, 
            author={Maximilian Herde and Bogdan Raonić and Tobias Rohner and Roger Käppeli and Roberto Molinaro and Emmanuel de Bézenac and Siddhartha Mishra},
            year={2024},
            eprint={2405.19101},
            archivePrefix={arXiv},
            primaryClass={cs.LG}
        }
        ```
        """
    )

demo.launch()
# demo.launch(server_name="0.0.0.0", server_port=7860)