app2.py

import streamlit as st
import torch
import random
import numpy as np
import yaml
from pathlib import Path
import tempfile
import traceback
import matplotlib.pyplot as plt
import plotly.graph_objects as go
from Prithvi import *  # Ensure this import includes your model and dataset classes
import xarray as xr
from aurora import Batch, Metadata
from aurora import Aurora, rollout
import logging
import matplotlib.pyplot as plt
import numpy as np
import cartopy.crs as ccrs
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Function to save uploaded files to temporary files and store paths in session_state
def save_uploaded_files(uploaded_files):
    if 'temp_file_paths' not in st.session_state:
        st.session_state.temp_file_paths = []
        for uploaded_file in uploaded_files:
            suffix = os.path.splitext(uploaded_file.name)[1]
            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
            temp_file.write(uploaded_file.read())
            temp_file.close()
            st.session_state.temp_file_paths.append(temp_file.name)
# Cached function to load dataset
@st.cache_resource
def load_dataset(file_paths):
    try:
        ds = xr.open_mfdataset(file_paths, combine='by_coords').load()
        return ds
    except Exception as e:
        st.error("Error loading dataset:")
        st.error(traceback.format_exc())
        return None

# Set page configuration
st.set_page_config(
    page_title="Weather Data Processor",
    layout="wide",
    initial_sidebar_state="expanded",
)



# Create a header with two columns: one for the title and one for the model selector
header_col1, header_col2 = st.columns([4, 1])  # Adjust the ratio as needed

with header_col1:
    st.title("🌦️ Weather & Climate Data Processor and Forecaster")

with header_col2:
    st.markdown("### Select a Model")
    selected_model = st.selectbox(
        "",
        options=["Aurora", "Climax", "Prithvi", "LSTM"],
        index=0,
        key="model_selector",
        help="Select the model you want to use for processing the data."
    )

st.write("---")  # Horizontal separator

# --- Layout: Two Columns ---
left_col, right_col = st.columns([1, 2])  # Adjust column ratios as needed

with left_col:
    st.header("🔧 Configuration")

    # --- Dynamic Configuration Based on Selected Model ---
    def get_model_configuration(model_name):
        if model_name == "Prithvi":
            st.subheader("Prithvi Model Configuration")

            # Prithvi-specific configuration inputs
            param1 = st.number_input("Prithvi Parameter 1", value=10, step=1)
            param2 = st.text_input("Prithvi Parameter 2", value="default_prithvi")
            # Add other Prithvi-specific parameters here

            config = {
                "param1": param1,
                "param2": param2,
                # Include other parameters as needed
            }

            # --- Prithvi-Specific File Uploads ---
            st.markdown("### Upload Data Files for Prithvi Model")

            # File uploader for surface data
            uploaded_surface_files = st.file_uploader(
                "Upload Surface Data Files",
                type=["nc", "netcdf"],
                accept_multiple_files=True,
                key="surface_uploader",
            )

            # File uploader for vertical data
            uploaded_vertical_files = st.file_uploader(
                "Upload Vertical Data Files",
                type=["nc", "netcdf"],
                accept_multiple_files=True,
                key="vertical_uploader",
            )

            # Handle Climatology Files
            st.markdown("### Upload Climatology Files (If Missing)")

            # Climatology files paths
            default_clim_dir = Path("Prithvi-WxC/examples/climatology")
            surf_in_scal_path = default_clim_dir / "musigma_surface.nc"
            vert_in_scal_path = default_clim_dir / "musigma_vertical.nc"
            surf_out_scal_path = default_clim_dir / "anomaly_variance_surface.nc"
            vert_out_scal_path = default_clim_dir / "anomaly_variance_vertical.nc"

            # Check if climatology files exist
            clim_files_exist = all(
                [
                    surf_in_scal_path.exists(),
                    vert_in_scal_path.exists(),
                    surf_out_scal_path.exists(),
                    vert_out_scal_path.exists(),
                ]
            )

            if not clim_files_exist:
                st.warning("Climatology files are missing.")
                uploaded_clim_surface = st.file_uploader(
                    "Upload Climatology Surface File",
                    type=["nc", "netcdf"],
                    key="clim_surface_uploader",
                )
                uploaded_clim_vertical = st.file_uploader(
                    "Upload Climatology Vertical File",
                    type=["nc", "netcdf"],
                    key="clim_vertical_uploader",
                )

                # Process uploaded climatology files
                if uploaded_clim_surface and uploaded_clim_vertical:
                    clim_temp_dir = tempfile.mkdtemp()
                    clim_surf_path = Path(clim_temp_dir) / uploaded_clim_surface.name
                    with open(clim_surf_path, "wb") as f:
                        f.write(uploaded_clim_surface.getbuffer())
                    clim_vert_path = Path(clim_temp_dir) / uploaded_clim_vertical.name
                    with open(clim_vert_path, "wb") as f:
                        f.write(uploaded_clim_vertical.getbuffer())
                    st.success("Climatology files uploaded and saved.")
                else:
                    st.warning("Please upload both climatology surface and vertical files.")
            else:
                clim_surf_path = surf_in_scal_path
                clim_vert_path = vert_in_scal_path

            # Optional: Upload config.yaml
            uploaded_config = st.file_uploader(
                "Upload config.yaml",
                type=["yaml", "yml"],
                key="config_uploader",
            )

            if uploaded_config:
                temp_config = tempfile.mktemp(suffix=".yaml")
                with open(temp_config, "wb") as f:
                    f.write(uploaded_config.getbuffer())
                config_path = Path(temp_config)
                st.success("Config.yaml uploaded and saved.")
            else:
                # Use default config.yaml path
                config_path = Path("Prithvi-WxC/examples/config.yaml")
                if not config_path.exists():
                    st.error("Default config.yaml not found. Please upload a config file.")
                    st.stop()

            # Optional: Upload model weights
            uploaded_weights = st.file_uploader(
                "Upload Model Weights (.pt)",
                type=["pt"],
                key="weights_uploader",
            )

            if uploaded_weights:
                temp_weights = tempfile.mktemp(suffix=".pt")
                with open(temp_weights, "wb") as f:
                    f.write(uploaded_weights.getbuffer())
                weights_path = Path(temp_weights)
                st.success("Model weights uploaded and saved.")
            else:
                # Use default weights path
                weights_path = Path("Prithvi-WxC/examples/weights/prithvi.wxc.2300m.v1.pt")
                if not weights_path.exists():
                    st.error("Default model weights not found. Please upload model weights.")
                    st.stop()

            return config, uploaded_surface_files, uploaded_vertical_files, clim_surf_path, clim_vert_path, config_path, weights_path

        else:
            # For other models, provide a simple file uploader
            st.subheader(f"{model_name} Model Data Upload")
            st.markdown("### Drag and Drop Your Data Files Here")
            uploaded_files = st.file_uploader(
                f"Upload Data Files for {model_name}",
                accept_multiple_files=True,
                key=f"{model_name.lower()}_uploader",
                type=["nc", "netcdf", "nc4"],
            )
            return uploaded_files

    # Retrieve model-specific configuration and files
    if selected_model == "Prithvi":
        config, uploaded_surface_files, uploaded_vertical_files, clim_surf_path, clim_vert_path, config_path, weights_path = get_model_configuration(selected_model)
    else:
        uploaded_files = get_model_configuration(selected_model)

    st.write("---")  # Horizontal separator

    # --- Run Inference Button ---
    if st.button("🚀 Run Inference"):
        with right_col:
            st.header("📈 Inference Progress & Visualization")

            # Initialize device
            try:
                torch.jit.enable_onednn_fusion(True)
                if torch.cuda.is_available():
                    device = torch.device("cuda")
                    st.write(f"Using device: **{torch.cuda.get_device_name()}**")
                    torch.backends.cudnn.benchmark = True
                    torch.backends.cudnn.deterministic = True
                else:
                    device = torch.device("cpu")
                    st.write("Using device: **CPU**")
            except Exception as e:
                st.error("Error initializing device:")
                st.error(traceback.format_exc())
                st.stop()

            # Set random seeds
            try:
                random.seed(42)
                if torch.cuda.is_available():
                    torch.cuda.manual_seed(42)
                torch.manual_seed(42)
                np.random.seed(42)
            except Exception as e:
                st.error("Error setting random seeds:")
                st.error(traceback.format_exc())
                st.stop()

            # # Define variables and parameters based on dataset type
            # if dataset_type == "MERRA2":
            #     surface_vars = [
            #         "EFLUX",
            #         "GWETROOT",
            #         "HFLUX",
            #         "LAI",
            #         "LWGAB",
            #         "LWGEM",
            #         "LWTUP",
            #         "PS",
            #         "QV2M",
            #         "SLP",
            #         "SWGNT",
            #         "SWTNT",
            #         "T2M",
            #         "TQI",
            #         "TQL",
            #         "TQV",
            #         "TS",
            #         "U10M",
            #         "V10M",
            #         "Z0M",
            #     ]
            #     static_surface_vars = ["FRACI", "FRLAND", "FROCEAN", "PHIS"]
            #     vertical_vars = ["CLOUD", "H", "OMEGA", "PL", "QI", "QL", "QV", "T", "U", "V"]
            #     levels = [
            #         34.0,
            #         39.0,
            #         41.0,
            #         43.0,
            #         44.0,
            #         45.0,
            #         48.0,
            #         51.0,
            #         53.0,
            #         56.0,
            #         63.0,
            #         68.0,
            #         71.0,
            #         72.0,
            #     ]
            # elif dataset_type == "GEOS5":
            #     # Define GEOS5 specific variables
            #     surface_vars = [
            #         "GEOS5_EFLUX",
            #         "GEOS5_GWETROOT",
            #         "GEOS5_HFLUX",
            #         "GEOS5_LAI",
            #         "GEOS5_LWGAB",
            #         "GEOS5_LWGEM",
            #         "GEOS5_LWTUP",
            #         "GEOS5_PS",
            #         "GEOS5_QV2M",
            #         "GEOS5_SLP",
            #         "GEOS5_SWGNT",
            #         "GEOS5_SWTNT",
            #         "GEOS5_T2M",
            #         "GEOS5_TQI",
            #         "GEOS5_TQL",
            #         "GEOS5_TQV",
            #         "GEOS5_TS",
            #         "GEOS5_U10M",
            #         "GEOS5_V10M",
            #         "GEOS5_Z0M",
            #     ]
            #     static_surface_vars = ["GEOS5_FRACI", "GEOS5_FRLAND", "GEOS5_FROCEAN", "GEOS5_PHIS"]
            #     vertical_vars = ["GEOS5_CLOUD", "GEOS5_H", "GEOS5_OMEGA", "GEOS5_PL", "GEOS5_QI", "GEOS5_QL", "GEOS5_QV", "GEOS5_T", "GEOS5_U", "GEOS5_V"]
            #     levels = [
            #         # Define levels specific to GEOS5 if different
            #         10.0,
            #         20.0,
            #         30.0,
            #         40.0,
            #         50.0,
            #         60.0,
            #         70.0,
            #         80.0,
            #     ]
            # else:
            #     st.error("Unsupported dataset type selected.")
            #     st.stop()

            padding = {"level": [0, 0], "lat": [0, -1], "lon": [0, 0]}

            residual = "climate"
            masking_mode = "local"
            decoder_shifting = True
            masking_ratio = 0.99

            positional_encoding = "fourier"

            # --- Initialize Dataset ---
            try:
                with st.spinner("Initializing dataset..."):
                    if selected_model == "Prithvi":
                        pass
                        # # Validate climatology files
                        # if not clim_files_exist and not (uploaded_clim_surface and uploaded_clim_vertical):
                        #     st.error("Climatology files are missing. Please upload both climatology surface and vertical files.")
                        #     st.stop()

                        # dataset = Merra2Dataset(
                        #     time_range=time_range,
                        #     lead_times=lead_times,
                        #     input_times=input_times,
                        #     data_path_surface=surf_dir,
                        #     data_path_vertical=vert_dir,
                        #     climatology_path_surface=clim_surf_path,
                        #     climatology_path_vertical=clim_vert_path,
                        #     surface_vars=surface_vars,
                        #     static_surface_vars=static_surface_vars,
                        #     vertical_vars=vertical_vars,
                        #     levels=levels,
                        #     positional_encoding=positional_encoding,
                        # )
                        # assert len(dataset) > 0, "There doesn't seem to be any valid data."
                    elif selected_model == "Aurora":
                        # TODO just temporary, replace this
                        if uploaded_files:
                            temp_file_paths = []  # List to store paths of temporary files
                            try:
                                # Save each uploaded file to a temporary file
                                save_uploaded_files(uploaded_files)
                                ds = load_dataset(st.session_state.temp_file_paths)

                                # Now, use xarray to open the multiple files
                                if ds:
                                    st.success("Files successfully loaded!")
                                    st.session_state.ds_subset = ds

                                    
                                    # print(ds)
                                    ds = ds.fillna(ds.mean())

                                    desired_levels = [50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000]

                                    # Ensure that the 'lev' dimension exists
                                    if 'lev' not in ds.dims:
                                        raise ValueError("The dataset does not contain a 'lev' (pressure level) dimension.")

                                    # Define the _prepare function
                                    def _prepare(x: np.ndarray, i: int) -> torch.Tensor:
                                        # Select previous and current time steps
                                        selected = x[[i - 6, i]]
                                        
                                        # Add a batch dimension
                                        selected = selected[None]
                                        
                                        # Ensure data is contiguous
                                        selected = selected.copy()
                                        
                                        # Convert to PyTorch tensor
                                        return torch.from_numpy(selected)

                                    # Adjust latitudes and longitudes
                                    lat = ds.lat.values * -1
                                    lon = ds.lon.values + 180

                                    # Subset the dataset to only include the desired pressure levels
                                    ds_subset = ds.sel(lev=desired_levels, method="nearest")

                                    # Verify that all desired levels are present
                                    present_levels = ds_subset.lev.values
                                    missing_levels = set(desired_levels) - set(present_levels)
                                    if missing_levels:
                                        raise ValueError(f"The following desired pressure levels are missing in the dataset: {missing_levels}")

                                    # Extract pressure levels after subsetting
                                    lev = ds_subset.lev.values  # Pressure levels in hPa

                                    # Prepare surface variables at 1000 hPa
                                    try:
                                        lev_index_1000 = np.where(lev == 1000)[0][0]
                                    except IndexError:
                                        raise ValueError("1000 hPa level not found in the 'lev' dimension after subsetting.")

                                    T_surface = ds_subset.T.isel(lev=lev_index_1000).compute()
                                    U_surface = ds_subset.U.isel(lev=lev_index_1000).compute()
                                    V_surface = ds_subset.V.isel(lev=lev_index_1000).compute()
                                    SLP = ds_subset.SLP.compute()

                                    # Reorder static variables (selecting the first time index to remove the time dimension)
                                    PHIS = ds_subset.PHIS.isel(time=0).compute()

                                    # Prepare atmospheric variables for the desired pressure levels excluding 1000 hPa
                                    atmos_levels = [int(level) for level in lev if level != 1000]

                                    T_atm = (ds_subset.T.sel(lev=atmos_levels)).compute()
                                    U_atm = (ds_subset.U.sel(lev=atmos_levels)).compute()
                                    V_atm = (ds_subset.V.sel(lev=atmos_levels)).compute()

                                    # Select time index
                                    num_times = ds_subset.time.size
                                    i = 6  # Adjust as needed (1 <= i < num_times)

                                    if i >= num_times or i < 1:
                                        raise IndexError("Time index i is out of bounds.")

                                    time_values = ds_subset.time.values
                                    current_time = np.datetime64(time_values[i]).astype('datetime64[s]').astype(datetime)

                                    # Prepare surface variables
                                    surf_vars = {
                                        "2t": _prepare(T_surface.values, i),   # Two-meter temperature
                                        "10u": _prepare(U_surface.values, i),  # Ten-meter eastward wind
                                        "10v": _prepare(V_surface.values, i),  # Ten-meter northward wind
                                        "msl": _prepare(SLP.values, i),        # Mean sea-level pressure
                                    }

                                    # Prepare static variables (now 2D tensors)
                                    static_vars = {
                                        "z": torch.from_numpy(PHIS.values.copy()),  # Geopotential (h, w)
                                        # Add 'lsm' and 'slt' if available and needed
                                    }

                                    # Prepare atmospheric variables
                                    atmos_vars = {
                                        "t": _prepare(T_atm.values, i),  # Temperature at desired levels
                                        "u": _prepare(U_atm.values, i),  # Eastward wind at desired levels
                                        "v": _prepare(V_atm.values, i),  # Southward wind at desired levels
                                    }

                                    # Define metadata
                                    metadata = Metadata(
                                        lat=torch.from_numpy(lat.copy()),
                                        lon=torch.from_numpy(lon.copy()),
                                        time=(current_time,),
                                        atmos_levels=tuple(atmos_levels),  # Only the desired atmospheric levels
                                    )

                                    # Create the Batch object
                                    batch = Batch(
                                        surf_vars=surf_vars,
                                        static_vars=static_vars,
                                        atmos_vars=atmos_vars,
                                        metadata=metadata
                                    ) # Display the dataset or perform further processing

                                    st.session_state['batch'] = batch

                            except Exception as e:
                                st.error(f"An error occurred: {e}")

                            # finally:
                            #     # Clean up: Remove temporary files
                            #     for path in temp_file_paths:
                            #         try:
                            #             os.remove(path)
                            #         except Exception as e:
                            #             st.warning(f"Could not delete temp file {path}: {e}")
                    else:
                        # For other models, implement their specific dataset initialization
                        # Placeholder: Replace with actual dataset initialization for other models
                        dataset = None  # Replace with actual dataset
                        st.warning("Dataset initialization for this model is not implemented yet.")
                        st.stop()
                st.success("Dataset initialized successfully.")
            except Exception as e:
                st.error("Error initializing dataset:")
                st.error(traceback.format_exc())
                st.stop()

            # --- Load Scalers ---
            try:
                with st.spinner("Loading scalers..."):
                    if selected_model == "Prithvi":
                        pass
                        # # Assuming the scaler paths are the same as climatology paths
                        # surf_in_scal_path = clim_surf_path
                        # vert_in_scal_path = clim_vert_path
                        # surf_out_scal_path = Path(clim_surf_path.parent) / "anomaly_variance_surface.nc"
                        # vert_out_scal_path = Path(clim_vert_path.parent) / "anomaly_variance_vertical.nc"

                        # # Check if output scaler files exist
                        # if not surf_out_scal_path.exists() or not vert_out_scal_path.exists():
                        #     st.error("Anomaly variance scaler files are missing.")
                        #     st.stop()

                        # in_mu, in_sig = input_scalers(
                        #     surface_vars,
                        #     vertical_vars,
                        #     levels,
                        #     surf_in_scal_path,
                        #     vert_in_scal_path,
                        # )

                        # output_sig = output_scalers(
                        #     surface_vars,
                        #     vertical_vars,
                        #     levels,
                        #     surf_out_scal_path,
                        #     vert_out_scal_path,
                        # )

                        # static_mu, static_sig = static_input_scalers(
                        #     surf_in_scal_path,
                        #     static_surface_vars,
                        # )
                    else:
                        # Load scalers for other models if applicable
                        # Placeholder: Replace with actual scaler loading for other models
                        in_mu, in_sig = None, None
                        output_sig = None
                        static_mu, static_sig = None, None
                st.success("Scalers loaded successfully.")
            except Exception as e:
                st.error("Error loading scalers:")
                st.error(traceback.format_exc())
                st.stop()

            # --- Load Configuration ---
            try:
                with st.spinner("Loading configuration..."):
                    if selected_model == "Prithvi":
                        with open(config_path, "r") as f:
                            config = yaml.safe_load(f)
                        # Validate config
                        required_params = [
                            "in_channels", "input_size_time", "in_channels_static",
                            "input_scalers_epsilon", "static_input_scalers_epsilon",
                            "n_lats_px", "n_lons_px", "patch_size_px",
                            "mask_unit_size_px", "embed_dim", "n_blocks_encoder",
                            "n_blocks_decoder", "mlp_multiplier", "n_heads",
                            "dropout", "drop_path", "parameter_dropout"
                        ]
                        missing_params = [param for param in required_params if param not in config.get("params", {})]
                        if missing_params:
                            st.error(f"Missing configuration parameters: {missing_params}")
                            st.stop()
                    else:
                        # Load configuration for other models if applicable
                        # Placeholder: Replace with actual configuration loading for other models
                        config = {}
                st.success("Configuration loaded successfully.")
            except Exception as e:
                st.error("Error loading configuration:")
                st.error(traceback.format_exc())
                st.stop()

            # --- Initialize the Model ---
            try:
                with st.spinner("Initializing model..."):
                    if selected_model == "Prithvi":
                        model = PrithviWxC(
                            in_channels=config["params"]["in_channels"],
                            input_size_time=config["params"]["input_size_time"],
                            in_channels_static=config["params"]["in_channels_static"],
                            input_scalers_mu=in_mu,
                            input_scalers_sigma=in_sig,
                            input_scalers_epsilon=config["params"]["input_scalers_epsilon"],
                            static_input_scalers_mu=static_mu,
                            static_input_scalers_sigma=static_sig,
                            static_input_scalers_epsilon=config["params"]["static_input_scalers_epsilon"],
                            output_scalers=output_sig**0.5,
                            n_lats_px=config["params"]["n_lats_px"],
                            n_lons_px=config["params"]["n_lons_px"],
                            patch_size_px=config["params"]["patch_size_px"],
                            mask_unit_size_px=config["params"]["mask_unit_size_px"],
                            mask_ratio_inputs=masking_ratio,
                            embed_dim=config["params"]["embed_dim"],
                            n_blocks_encoder=config["params"]["n_blocks_encoder"],
                            n_blocks_decoder=config["params"]["n_blocks_decoder"],
                            mlp_multiplier=config["params"]["mlp_multiplier"],
                            n_heads=config["params"]["n_heads"],
                            dropout=config["params"]["dropout"],
                            drop_path=config["params"]["drop_path"],
                            parameter_dropout=config["params"]["parameter_dropout"],
                            residual=residual,
                            masking_mode=masking_mode,
                            decoder_shifting=decoder_shifting,
                            positional_encoding=positional_encoding,
                            checkpoint_encoder=[],
                            checkpoint_decoder=[],
                        )
                    elif selected_model == "Aurora":
                        pass
                        
                    else:
                        
                        # Initialize other models here
                        # Placeholder: Replace with actual model initialization for other models
                        model = None
                        st.warning("Model initialization for this model is not implemented yet.")
                        st.stop()
                    # model.to(device)
                st.success("Model initialized successfully.")
            except Exception as e:
                st.error("Error initializing model:")
                st.error(traceback.format_exc())
                st.stop()

            # --- Load Model Weights ---
            try:
                with st.spinner("Loading model weights..."):
                    if selected_model == "Prithvi":
                        state_dict = torch.load(weights_path, map_location=device)
                        if "model_state" in state_dict:
                            state_dict = state_dict["model_state"]
                        model.load_state_dict(state_dict, strict=True)
                        model.to(device)
                    else:
                        # Load weights for other models if applicable
                        # Placeholder: Replace with actual weight loading for other models
                        pass
                st.success("Model weights loaded successfully.")
            except Exception as e:
                st.error("Error loading model weights:")
                st.error(traceback.format_exc())
                st.stop()

            # --- Prepare Data Batch ---
            try:
                with st.spinner("Preparing data batch..."):
                    if selected_model == "Prithvi":
                        data = next(iter(dataset))
                        batch = preproc([data], padding)
                        for k, v in batch.items():
                            if isinstance(v, torch.Tensor):
                                batch[k] = v.to(device)
                    elif selected_model == "Aurora":
                        batch = batch.regrid(res=0.25)
                        
                    else:
                        # Prepare data batch for other models
                        # Placeholder: Replace with actual data preparation for other models
                        batch = None
                st.success("Data batch prepared successfully.")
            except Exception as e:
                st.error("Error preparing data batch:")
                st.error(traceback.format_exc())
                st.stop()

            # --- Run Inference ---
            try:
                with st.spinner("Running model inference..."):
                    if selected_model == "Prithvi":
                        model.eval()
                        with torch.no_grad():
                            out = model(batch)
                    elif selected_model == "Aurora":
                        
                        model = Aurora(use_lora=False)
                        # model = Aurora()
                        model.load_checkpoint("microsoft/aurora", "aurora-0.25-pretrained.ckpt")
                        # model.load_checkpoint("microsoft/aurora", "aurora-0.25-pretrained.ckpt")

                        model.eval()
                        # model = model.to("cuda")  # Uncomment if using a GPU

                        with torch.inference_mode():
                            out = [pred.to("cpu") for pred in rollout(model, batch, steps=2)]

                        model = model.to("cpu")
                        st.session_state.model = model
                    else:
                        # Run inference for other models
                        # Placeholder: Replace with actual inference code for other models
                        out = torch.randn(1, 10, 180, 360)  # Dummy tensor
                st.success("Model inference completed successfully.")
                st.session_state['out'] = out
            except Exception as e:
                st.error("Error during model inference:")
                st.error(traceback.format_exc())
                st.stop()

            # --- Visualization Settings ---
            st.markdown("## 📊 Visualization Settings")

            if 'out' in st.session_state and 'batch' in st.session_state and selected_model == "Prithvi":
                # Display the shape of the output tensor
                out_tensor = st.session_state['out']
                st.write(f"**Output tensor shape:** {out_tensor.shape}")

                # Ensure the output tensor has at least 4 dimensions (batch, variables, lat, lon)
                if out_tensor.ndim < 4:
                    st.error("The output tensor does not have the expected number of dimensions (batch, variables, lat, lon).")
                    st.stop()

                # Get the number of variables
                num_variables = out_tensor.shape[1]

                # Define variable names (update with your actual variable names)
                variable_names = [f"Variable_{i}" for i in range(num_variables)]

                # Visualization settings
                col1, col2 = st.columns(2)

                with col1:
                    # Select variable to plot
                    selected_variable_name = st.selectbox(
                        "Select Variable to Plot",
                        options=variable_names,
                        index=0,
                        help="Choose the variable you want to visualize."
                    )

                    # Select plot type
                    plot_type = st.selectbox(
                        "Select Plot Type",
                        options=["Contour", "Heatmap"],
                        index=0,
                        help="Choose the type of plot to display."
                    )

                with col2:
                    # Select color map
                    cmap = st.selectbox(
                        "Select Color Map",
                        options=plt.colormaps(),
                        index=plt.colormaps().index("viridis"),
                        help="Choose the color map for the plot."
                    )

                    # Set number of levels (for contour plot)
                    if plot_type == "Contour":
                        num_levels = st.slider(
                            "Number of Contour Levels",
                            min_value=5,
                            max_value=100,
                            value=20,
                            step=5,
                            help="Set the number of contour levels."
                        )
                    else:
                        num_levels = None

                # Find the index based on the selected name
                variable_index = variable_names.index(selected_variable_name)

                # Extract the selected variable
                selected_variable = out_tensor[0, variable_index].cpu().numpy()

                # Generate latitude and longitude arrays
                lat = np.linspace(-90, 90, selected_variable.shape[0])
                lon = np.linspace(-180, 180, selected_variable.shape[1])
                X, Y = np.meshgrid(lon, lat)

                # Plot the selected variable
                st.markdown(f"### Plot of {selected_variable_name}")

                # Matplotlib figure
                fig, ax = plt.subplots(figsize=(10, 6))

                if plot_type == "Contour":
                    # Generate the contour plot
                    contour = ax.contourf(X, Y, selected_variable, levels=num_levels, cmap=cmap)
                elif plot_type == "Heatmap":
                    # Generate the heatmap
                    contour = ax.imshow(selected_variable, extent=[-180, 180, -90, 90], cmap=cmap, origin='lower', aspect='auto')

                # Add a color bar
                cbar = plt.colorbar(contour, ax=ax)
                cbar.set_label(f'{selected_variable_name}', fontsize=12)

                # Set aspect ratio and labels
                ax.set_xlabel("Longitude", fontsize=12)
                ax.set_ylabel("Latitude", fontsize=12)
                ax.set_title(f"{selected_variable_name}", fontsize=14)

                # Display the plot in Streamlit
                st.pyplot(fig)

                # Optional: Provide interactive Plotly plot
                st.markdown("#### Interactive Plot")
                if plot_type == "Contour":
                    fig_plotly = go.Figure(data=go.Contour(
                        z=selected_variable,
                        x=lon,
                        y=lat,
                        colorscale=cmap,
                        contours=dict(
                            coloring='fill',
                            showlabels=True,
                            labelfont=dict(size=12, color='white'),
                            ncontours=num_levels
                        )
                    ))
                elif plot_type == "Heatmap":
                    fig_plotly = go.Figure(data=go.Heatmap(
                        z=selected_variable,
                        x=lon,
                        y=lat,
                        colorscale=cmap
                    ))

                fig_plotly.update_layout(
                    xaxis_title="Longitude",
                    yaxis_title="Latitude",
                    autosize=False,
                    width=800,
                    height=600,
                )

                st.plotly_chart(fig_plotly)

            elif 'out' in st.session_state and selected_model == "Aurora" and st.session_state['out'] is not None:
                preds = st.session_state['out']
                ds_subset = st.session_state.get('ds_subset', None)
                batch = st.session_state.get('batch', None)

                # **Determine Available Levels**
                # For example, let's assume levels range from 0 to max_level_index
                # You need to replace 'max_level_index' with the actual maximum level index in your data
                try:
                    # Assuming 'lev' dimension exists and is 1D
                    levels = preds[0].atmos_vars["t"].shape[2]  # Adjust based on your data structure
                    level_indices = list(range(levels))
                except Exception as e:
                    st.error("Error determining available levels:")
                    st.error(traceback.format_exc())
                    levels = None  # Set to None if levels cannot be determined

                if levels is not None:
                    # **Add a Slider for Level Selection**
                    selected_level = st.slider(
                        'Select Level',
                        min_value=0,
                        max_value=levels - 1,
                        value=11,  # Default level index
                        step=1,
                        help="Select the vertical level for plotting."
                    )

                    # Loop through predictions and ground truths
                    for idx in range(len(preds)):
                        pred = preds[idx]
                        pred_time = pred.metadata.time[0]

                        # Display prediction time
                        st.write(f"### Prediction Time: {pred_time}")

                        # **Extract Data at Selected Level**
                        try:
                            # Update indices with the selected_level
                            pred_data = pred.atmos_vars["t"][0][0][selected_level].numpy() - 273.15
                            truth_data = ds_subset.T.isel(lev=selected_level)[idx].values - 273.15
                            
                        except Exception as e:
                            st.error("Error extracting data for plotting:")
                            st.error(traceback.format_exc())
                            continue

                        # Extract latitude and longitude
                        try:
                            lat = np.array(pred.metadata.lat)  # Assuming 'lat' is 1D
                            lon = np.array(pred.metadata.lon)  # Assuming 'lon' is 1D
                        except Exception as e:
                            st.error("Error extracting latitude and longitude:")
                            st.error(traceback.format_exc())
                            continue

                        # Create a meshgrid for plotting
                        lon_grid, lat_grid = np.meshgrid(lon, lat)

                        # Create a Matplotlib figure with Cartopy projection
                        fig, axes = plt.subplots(
                            1, 3, figsize=(18, 6),
                            subplot_kw={'projection': ccrs.PlateCarree()}
                        )

                        # **Ground Truth Plot**
                        im1 = axes[0].imshow(
                            truth_data,
                            extent=[lon.min(), lon.max(), lat.min(), lat.max()],
                            origin='lower',
                            cmap='coolwarm',
                            transform=ccrs.PlateCarree()
                        )
                        axes[0].set_title(f"Ground Truth at Level {selected_level} - {pred_time}")
                        axes[0].set_xlabel('Longitude')
                        axes[0].set_ylabel('Latitude')
                        plt.colorbar(im1, ax=axes[0], orientation='horizontal', pad=0.05)

                        # **Prediction Plot**
                        im2 = axes[1].imshow(
                            pred_data,
                            extent=[lon.min(), lon.max(), lat.min(), lat.max()],
                            origin='lower',
                            cmap='coolwarm',
                            transform=ccrs.PlateCarree()
                        )
                        axes[1].set_title(f"Prediction at Level {selected_level} - {pred_time}")
                        axes[1].set_xlabel('Longitude')
                        axes[1].set_ylabel('Latitude')
                        plt.colorbar(im2, ax=axes[1], orientation='horizontal', pad=0.05)

                        plt.tight_layout()

                        # Display the plot in Streamlit
                        st.pyplot(fig)
                else:
                    st.error("Could not determine the available levels in the data.")


            else:
                st.warning("No output available to display or visualization is not implemented for this model.")

    # --- End of Inference Button ---
    else:
        with right_col:
            st.header("🖥️ Visualization & Progress")
            st.info("Awaiting inference to display results.")