app.py

from fluid import * 
from gradio_litmodel3d import LitModel3D
import math, zipfile, tempfile
import numpy as np
import plotly.graph_objects as go
import gradio as gr
import cudaq

# Helper: snap grid size to nearest power of two within bounds
def _nearest_pow2(n: int, min_pow: int = 2**7, max_pow: int = 2**12) -> int:
    if n <= 0:
        return min_pow
    # nearest power of two
    p = 1 << (int(round(math.log2(max(n, 1)))))
    # clamp to bounds and adjust if out of range
    if p < min_pow:
        p = min_pow
    if p > max_pow:
        # choose the closer between max_pow and previous power
        p = max_pow
    return p

# New update function: also updates the slider value after snapping
def update_grid_and_qubit_info(grid_size):
    snapped = _nearest_pow2(int(grid_size))
    num_reg_qubits = int(math.log2(snapped))
    total_qubits = 2 * num_reg_qubits + 3

    x = np.array([128, 256, 512, 1024, 2048, 4096])
    y = np.log2(x).astype(int)
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=x, y=y, mode='lines', name='Qubits/Direction'))
    fig.add_trace(go.Scatter(x=[snapped], y=[num_reg_qubits], mode='markers', marker=dict(size=12, color='red'), name='Current Selection'))
    fig.update_layout(
        xaxis_title="Grid Size (Points/Direction)",
        yaxis_title="Qubits/Direction",
        width=400,
        height=300
    )

    total_qubits_display = f"Total Qubits: {total_qubits}"
    warn_pow2 = " (snapped to nearest power of two)" if snapped != int(grid_size) else ""
    warning = ("⚠️ Warning: Grid sizes > 1024 may exceed simulator/memory limits!" if snapped > 1024 else "") + warn_pow2
    recommended_time_steps = num_reg_qubits * 200
    recommended_display = f"Recommended time steps: {recommended_time_steps}"

    return gr.update(value=snapped), fig, total_qubits_display, warning, recommended_display

# Modified interface functions to take num_reg_qubits_input
def qlbm_gradio_interface(grid_size_input: int, time_steps_input: int, distribution_type_param: str, velocity_field_param: str, vx_param: float, vy_param: float, boundary_condition_param: str):
    snapped_grid = _nearest_pow2(int(grid_size_input))
    if snapped_grid != int(grid_size_input):
        gr.Warning(f"Grid size {grid_size_input} is not a power of two. Using {snapped_grid}.")
    num_reg_qubits_val = int(math.log2(snapped_grid))
    grid_size_val = snapped_grid
    time_steps_val = int(time_steps_input)
    vx_val = float(vx_param)
    vy_val = float(vy_param)

    print(f"Gradio Interface: Qubits/Direction={num_reg_qubits_val}, Grid Size={grid_size_val}, T={time_steps_val}, Distribution={distribution_type_param}, Velocity Field={velocity_field_param}, vx={vx_val}, vy={vy_val}, Boundary={boundary_condition_param}")

    plot_fig, plotly_json_frames = simulate_qlbm_and_animate( # Modified to unpack two return values
        num_reg_qubits=num_reg_qubits_val,
        T=time_steps_val,
        distribution_type=distribution_type_param,
        velocity_field=velocity_field_param,
        vx_input=vx_val,
        vy_input=vy_val,
        boundary_condition=boundary_condition_param
    )

    if plot_fig is None:
        gr.Warning("Simulation or plotting failed. Please check console for errors.")
        return None, None # Modified return
    return plot_fig, plotly_json_frames # Modified return

# New functions for downloading Plotly objects
def download_plot_data(plotly_json_frames):
    if not plotly_json_frames:
        gr.Warning("No data to download.")
        return None

    zip_file_path = tempfile.NamedTemporaryFile(suffix=".zip", delete=False).name
    with zipfile.ZipFile(zip_file_path, 'w', zipfile.ZIP_DEFLATED) as zf:
        for i, json_str in enumerate(plotly_json_frames):
            zf.writestr(f"frame_{i}.json", json_str)
    return zip_file_path

# Modified update functions to take grid_size (which is num_reg_qubits_input now)
def update_qubit_info(grid_size):
    num_reg_qubits = int(math.log2(grid_size))
    total_qubits = 2 * num_reg_qubits + 3

    x = np.array([128, 256, 512, 1024, 2048, 4096])
    y = np.log2(x).astype(int)
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=x, y=y, mode='lines', name='Qubits/Direction'))
    fig.add_trace(go.Scatter(x=[grid_size], y=[num_reg_qubits], mode='markers', marker=dict(size=12, color='red'), name='Current Selection'))
    fig.update_layout(
        xaxis_title="Grid Size (Points/Direction)",
        yaxis_title="Qubits/Direction",
        width=400,
        height=300
    )

    total_qubits_display = f"Total Qubits: {total_qubits}"
    warning = "⚠️ Warning: Grid sizes > 1024 may exceed simulator/memory limits!" if grid_size > 1024 else ""
    recommended_time_steps = num_reg_qubits * 200
    recommended_display = f"Recommended time steps: {recommended_time_steps}"

    return fig, total_qubits_display, warning, recommended_display

# Modified example functions to set grid_size
def set_sinusoidal_example():
    return (
        gr.update(value=256),  # grid_size = 256 (corresponds to 8 qubits)
        gr.update(value=1600),  # time_steps
        gr.update(value="Sinusoidal"),  # distribution_type
        gr.update(value="Uniform"),  # velocity_field
        gr.update(value=0.2),  # vx
        gr.update(value=0.15),  # vy
        gr.update(value="Periodic"),  # boundary_condition
    )

def set_gaussian_example():
    return (
        gr.update(value=256),  # grid_size = 256 (corresponds to 8 qubits)
        gr.update(value=1600),  # time_steps
        gr.update(value="Gaussian"),  # distribution_type
        gr.update(value="Uniform"),  # velocity_field
        gr.update(value=0.2),  # vx
        gr.update(value=0.15),  # vy
        gr.update(value="Periodic"),  # boundary_condition
    )

# Gradio interface for Fluid Dynamics - 2D only
with gr.Blocks(theme=gr.themes.Soft(), title="QLBM Fluid Simulation (2D)") as demo:
    with gr.Tabs():
        with gr.TabItem("Introduction"):
            gr.Markdown(
                """
                # Quantum Lattice Boltzmann (QLBM)
                This app runs a 2D quantum-inspired Lattice Boltzmann Method using a D2Q5 model (5 discrete velocity directions) to evolve a scalar density field with uniform advection and periodic boundaries.

                What is simulated
                - D2Q5 advection–diffusion on a square grid (periodic BCs only).
                - Uniform velocity (Vx, Vy) applied to all cells.
                - State is evolved via CUDA-Q kernels; results are downsampled and visualized as 3D surfaces with a time slider.

                How it works (from the implementation)
                - Grid size N = 2^q (q = qubits per spatial dimension). Total qubits = 2*q + 3 direction qubits.
                - Per timestep: a Householder-based prep_op initializes direction amplitudes, then conditional shifts stream along x and y (lshift/rshift) controlled by direction bits, followed by an unprep.
                - Up to 40 uniform time samples between 0..T are saved for visualization.

                Visualization pipeline
                - Only the real component of the state is saved and plotted; the initialized field is normalized before evolution.
                - Data is saved as .npy in a temporary folder, downsampled by a factor of 2^5 (32). For very small grids, the factor is reduced to keep at least 1×1.
                - Plotly renders one surface per saved frame; a slider toggles visibility. You can download per-frame Plotly JSON for offline analysis.

                Parameters you control
                - Grid Size/Direction: N = 2^q. The UI snaps any input to the nearest power of two in [128, 4096]. Larger N increases memory and runtime.
                - Time Steps (T): total evolution steps; up to 40 frames are sampled in [0, T].
                - Initial Distribution: Sinusoidal or Gaussian.
                - Velocity Field: Uniform (set Vx, Vy).
                - Boundary Condition: Periodic.

                Practical notes
                - Recommended time steps ≈ q * 200.
                - N > 1024 may exceed available memory/time.
                - Requires a CUDA-capable GPU and CUDA-Q runtime; CPU fallback will be slow or may fail.

                Workflow
                1) Review example initial distributions.
                2) Choose grid size N and time steps T.
                3) Set Vx and Vy; keep Periodic BC.
                4) Run the simulation; use the slider to browse frames.
                5) Optionally download the plotted frames as JSON for offline analysis.
                """
            )

        with gr.TabItem("Simulation"):
            with gr.Row(): # Main row for top section
                with gr.Column(scale=1): # Column for left-side controls
                    gr.Markdown("## Initial Distribution Examples")
                    with gr.Row():
                        with gr.Column(scale=1):
                            example1 = LitModel3D("Placeholder_Images/sinusoidal.stl", label="Sinusoidal")
                            sinusoidal_btn_2d = gr.Button("Sinusoidal")
                        with gr.Column(scale=1):
                            example2 = LitModel3D("Placeholder_Images/gaussian.stl", label="Gaussian")
                            gaussian_btn_2d = gr.Button("Gaussian")

                    gr.Markdown("## Simulation Parameters")
                    num_reg_qubits_input_2d = gr.Slider(
                        minimum=2**7, maximum=2**12, # Grid size values
                        value=2**8, step=1, # use step=1 for broader Gradio compatibility
                        label="Grid Size/Direction"
                    )
                    time_steps_slider_2d = gr.Slider(minimum=0, maximum=4000, value=1600, step=10, label="Time Steps")
                    qubit_plot_2d = gr.Plot(label="Qubits vs. Grid Size")
                    total_qubits_display_2d = gr.Markdown("Total Qubits: 19")
                    warning_display_2d = gr.Markdown("")
                    recommended_time_steps_display_2d = gr.Markdown("Recommended time steps: 1600")

                with gr.Column(scale=2): # Column for the main plot
                    qlbm_interactive_plot_2d = gr.Plot(label="QLBM")
                    download_button_2d = gr.DownloadButton(label="Download Plot Data (JSON)", visible=False) # New download button
                    plotly_json_frames_state_2d = gr.State([]) # New state to store Plotly JSON frames

            with gr.Row(): # Row for bottom section
                with gr.Column(scale=1):
                    gr.Markdown("## Initialization (Uniform)")
                    selected_velocity_field_2d = gr.State("Uniform")
                    vx_slider_2d = gr.Slider(minimum=-0.3, maximum=0.3, value=0.2, step=0.01, label="V_x", visible=True)
                    vy_slider_2d = gr.Slider(minimum=-0.3, maximum=0.3, value=0.15, step=0.01, label="V_y", visible=True)
                    distribution_type_input_2d = gr.Radio(choices=["Gaussian", "Sinusoidal"], value="Sinusoidal", label="Initial Distribution Type")

                with gr.Column(scale=1):
                    gr.Markdown("## Boundary Conditions")
                    boundary_condition_input_2d = gr.Radio(choices=["Periodic"], value="Periodic", label="Boundary Condition")

            run_qlbm_btn_2d = gr.Button("Run Simulation", variant="primary") # Button below the sections

            qlbm_inputs_list_2d = [
                num_reg_qubits_input_2d, time_steps_slider_2d, distribution_type_input_2d,
                selected_velocity_field_2d, vx_slider_2d, vy_slider_2d, boundary_condition_input_2d
            ]
            run_qlbm_btn_2d.click(
                fn=qlbm_gradio_interface,
                inputs=qlbm_inputs_list_2d,
                outputs=[qlbm_interactive_plot_2d, plotly_json_frames_state_2d] # Modified output
            ).then(
                lambda: gr.update(visible=True), # Make download button visible after simulation
                outputs=[download_button_2d]
            )
            download_button_2d.click(
                fn=download_plot_data,
                inputs=[plotly_json_frames_state_2d],
                outputs=[download_button_2d]
            )
            num_reg_qubits_input_2d.change(
                fn=update_grid_and_qubit_info,
                inputs=num_reg_qubits_input_2d,
                outputs=[num_reg_qubits_input_2d, qubit_plot_2d, total_qubits_display_2d, warning_display_2d, recommended_time_steps_display_2d]
            )
            sinusoidal_btn_2d.click(
                fn=set_sinusoidal_example,
                inputs=[],
                outputs=[num_reg_qubits_input_2d, time_steps_slider_2d, distribution_type_input_2d, selected_velocity_field_2d, vx_slider_2d, vy_slider_2d, boundary_condition_input_2d]
            )
            gaussian_btn_2d.click(
                fn=set_gaussian_example,
                inputs=[],
                outputs=[num_reg_qubits_input_2d, time_steps_slider_2d, distribution_type_input_2d, selected_velocity_field_2d, vx_slider_2d, vy_slider_2d, boundary_condition_input_2d]
            )

if __name__ == "__main__":
    try:
        cudaq.set_target('nvidia', option='fp64')
        print(f"CUDA-Q Target successfully set to: {cudaq.get_target().name}")
    except Exception as e_target:
        print(f"Warning: Could not set CUDA-Q target to 'nvidia'. Error: {e_target}")
        print(f"Current CUDA-Q Target: {cudaq.get_target().name}. Performance may be affected.")

    # Force local serving; disable any implicit sharing/tunneling
    import os
    os.environ["GRADIO_SHARE"] = "0"
    os.environ["GRADIO_ANALYTICS_ENABLED"] = "0"
    os.environ["GRADIO_LAUNCH_BROWSER"] = "0"

    # Bind to all interfaces so the host can reach it; no share
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False, inbrowser=False, show_error=True)