import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import tempfile
import gradio as gr

def solve_and_animate_2d(Lx: float,
                         Ly: float,
                         t_max: float,
                         Gamma: float = 0.1,
                         Nx: int = 50,
                         Ny: int = 50,
                         initial: str = "gaussian",
                         bc: str = "dirichlet",
                         frame_skip: int = 1):
    """
    Solve the 2D heat equation u_t = Gamma*(u_xx + u_yy) and return a GIF animation.
    Initial conditions: {"gaussian", "random", "sinusoidal", "step"}
    Boundary conditions: {"dirichlet", "neumann", "periodic"}
    """
    # Spatial grid
    x = np.linspace(0, Lx, Nx)
    y = np.linspace(0, Ly, Ny)
    dx, dy = x[1] - x[0], y[1] - y[0]
    if dx == 0 or dy == 0:
        raise ValueError("Nx and Ny must be > 1.")

    # Time stepping for stability
    dt = 1.0 / (2 * Gamma * (1/dx**2 + 1/dy**2))
    Nt = int(np.ceil(t_max / dt)) + 1
    rx, ry = Gamma * dt / dx**2, Gamma * dt / dy**2

    # Initial condition
    X, Y = np.meshgrid(x, y, indexing='ij')
    if initial == "gaussian":
        u = np.exp(-(((X - Lx/2)**2 + (Y - Ly/2)**2) / (2*(Lx/10)**2)))
    elif initial == "random":
        u = np.random.rand(Nx, Ny)
    elif initial == "sinusoidal":
        kx, ky = 2 * np.pi / Lx, 2 * np.pi / Ly
        u = np.sin(kx * X) * np.sin(ky * Y)
    elif initial == "step":
        u = np.where((X < Lx/2) & (Y < Ly/2), 1.0, 0.0)
    else:
        raise ValueError(f"Unknown initial condition: {initial}")

    # Storage for solution
    U = np.zeros((Nt, Nx, Ny))
    U[0] = u.copy()

    # Time-stepping loop
    for n in range(1, Nt):
        un = u.copy()
        # Interior update
        u[1:-1, 1:-1] = (
            un[1:-1, 1:-1]
            + rx * (un[2:, 1:-1] - 2 * un[1:-1, 1:-1] + un[:-2, 1:-1])
            + ry * (un[1:-1, 2:] - 2 * un[1:-1, 1:-1] + un[1:-1, :-2])
        )
        # Boundary conditions
        if bc == "dirichlet":
            u[0, :] = u[-1, :] = u[:, 0] = u[:, -1] = 0.0
        elif bc == "neumann":
            u[0, :] = u[1, :]
            u[-1, :] = u[-2, :]
            u[:, 0] = u[:, 1]
            u[:, -1] = u[:, -2]
        elif bc == "periodic":
            u[0, :] = un[-2, :]
            u[-1, :] = un[1, :]
            u[:, 0] = un[:, -2]
            u[:, -1] = un[:, 1]
        else:
            raise ValueError(f"Unknown bc: {bc}")
        U[n] = u.copy()

    # Create animation with Matplotlib
    fig, ax = plt.subplots()
    im = ax.imshow(U[0].T, cmap='viridis', origin='lower', extent=[0, Lx, 0, Ly], vmin=U.min(), vmax=U.max())
    ax.set_title(f"2D Heat Eq — init={initial}, bc={bc}, Gamma={Gamma:.2f}")
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    plt.colorbar(im, ax=ax, label="u")

    # Animation update function
    def update(frame):
        im.set_data(U[frame].T)
        return [im]

    # Subsample frames
    idx = list(range(0, Nt, frame_skip))
    if idx[-1] != Nt - 1:
        idx.append(Nt - 1)

    # Generate animation
    ani = FuncAnimation(fig, update, frames=idx, blit=True)

    # Save as GIF
    with tempfile.NamedTemporaryFile(suffix='.gif', delete=False) as tmpfile:
        ani.save(tmpfile.name, writer='pillow', fps=30)
        gif_path = tmpfile.name

    plt.close(fig)
    return gif_path

def gradio_interface(lx, ly, t_max, gamma, nx, ny, initial, bc, frame_skip):
    nx, ny, frame_skip = int(nx), int(ny), int(frame_skip)
    return solve_and_animate_2d(
        Lx=lx, Ly=ly, t_max=t_max, Gamma=gamma, Nx=nx, Ny=ny,
        initial=initial, bc=bc, frame_skip=frame_skip
    )

with gr.Blocks(theme=gr.themes.Soft(), title="2D Heat Simulator") as demo:
    gr.Markdown("# ♨️ 2D Heat Equation Simulator\nAdjust parameters and run the simulation.")
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("## Domain & Grid")
            lx_slider = gr.Slider(0.1, 5.0, 1.0, 0.1, label="Lx")
            ly_slider = gr.Slider(0.1, 5.0, 1.0, 0.1, label="Ly")
            nx_slider = gr.Slider(3, 200, 50, 1, label="Nx")
            ny_slider = gr.Slider(3, 200, 50, 1, label="Ny")

            gr.Markdown("## Simulation")
            t_slider = gr.Slider(0.01, 5.0, 0.5, 0.01, label="t_max")
            gamma_slider = gr.Slider(0.001, 1.0, 0.1, 0.001, label="Gamma")

            gr.Markdown("## Conditions")
            initial_dropdown = gr.Dropdown(
                ["gaussian", "random", "sinusoidal", "step"], "gaussian", label="Initial"
            )
            bc_dropdown = gr.Dropdown(
                ["dirichlet", "neumann", "periodic"], "dirichlet", label="Boundary"
            )

            gr.Markdown("## Animation")
            frame_skip_slider = gr.Slider(1, 50, 5, 1, label="Frame Skip")
            run_btn = gr.Button("Run Simulation", variant="primary")
        with gr.Column(scale=3):
            plot_output = gr.Image(label="Heatmap Animation")
    inputs_list = [lx_slider, ly_slider, t_slider, gamma_slider,
                   nx_slider, ny_slider, initial_dropdown, bc_dropdown, frame_skip_slider]
    run_btn.click(fn=gradio_interface, inputs=inputs_list, outputs=plot_output)
    gr.Examples(
        examples=[
            [1.0, 1.0, 0.5, 0.1, 50, 50, "gaussian", "dirichlet", 5],
            [2.0, 1.0, 1.0, 0.05, 60, 30, "sinusoidal", "periodic", 10],
            [1.0, 1.0, 0.2, 0.2, 80, 80, "step", "neumann", 2],
        ],
        inputs=inputs_list,
        outputs=[plot_output],
        fn=gradio_interface
    )

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