import numpy as np
import matplotlib
matplotlib.use("Agg")

import matplotlib.pyplot as plt
import tempfile

import gradio as gr

from bpm.refractive_index import (
    generate_waveguide_n_r2,
    generate_MMI_n_r2
)
from bpm.mode_solver import slab_mode_source
from bpm.core import run_bpm
from bpm.pml import generate_sigma_x

def run_waveguide(w, l, L_bend, n_WG, wavelength, ind_m):
    # --- Simulation setup (Waveguide S-bend) ---
    domain_size = 50.0
    z_total     = 500.0
    Nx, Nz      = 256, 2000

    x = np.linspace(-domain_size/2, domain_size/2, Nx)
    z = np.linspace(0, z_total, Nz)
    n0 = 1.0

    # Refractive index map
    n_r2 = generate_waveguide_n_r2(x, z, l, L_bend, w, n_WG, n0)

    # Launch slab mode
    E0 = slab_mode_source(x, w, n_WG, n0, wavelength, ind_m, x0=0)
    E  = np.zeros((Nx, Nz), dtype=np.complex128)
    E[:, 0] = E0

    # PML + BPM
    dx      = domain_size / Nx
    dz      = z[1] - z[0]
    sigma_x = generate_sigma_x(x, dx, wavelength, domain_size)
    E_out   = run_bpm(E, n_r2, x, z, dx, dz, n0, sigma_x, wavelength)

    # Plot
    fig, ax = plt.subplots(figsize=(8,6))
    im = ax.imshow(
        np.abs(E_out)**2,
        extent=[x[0], x[-1], z[0], z[-1]],
        origin='lower',
        aspect='auto',
        cmap='inferno'
    )
    ax.set_title("Waveguide S-bend BPM Propagation")
    ax.set_xlabel("x (µm)")
    ax.set_ylabel("z (µm)")
    fig.colorbar(im, ax=ax, label="Intensity")

    # Save data for download
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".npz")
    np.savez(tmp.name, E_out=E_out, x=x, z=z)
    tmp.close()

    return fig, tmp.name

def run_mmi(z_start, L_MMI, w_MMI, w_wg, d, n_WG, wavelength, ind_m):
    # --- Simulation setup (MMI Splitter) ---
    domain_size = 50.0
    z_total     = 250.0
    Nx, Nz      = 256, 1024

    x = np.linspace(-domain_size/2, domain_size/2, Nx)
    z = np.linspace(0, z_total, Nz)
    n0 = 1.0

    # For simplicity, use the same core and MMI index
    n_MMI = n_WG

    # Refractive index map
    n_r2 = generate_MMI_n_r2(
        x, z,
        z_MMI_start = z_start,
        L_MMI        = L_MMI,
        w_MMI        = w_MMI,
        w_wg         = w_wg,
        d            = d,
        n_WG         = n_WG,
        n_MMI        = n_MMI,
        n0           = n0
    )

    # Launch slab mode in left waveguide
    E0 = slab_mode_source(x, w_wg, n_WG, n0, wavelength, ind_m, x0=-d/2)
    E  = np.zeros((Nx, Nz), dtype=np.complex128)
    E[:,0] = E0

    # PML + BPM
    dx      = domain_size / Nx
    dz      = z[1] - z[0]
    sigma_x = generate_sigma_x(x, dx, wavelength, domain_size)
    E_out   = run_bpm(E, n_r2, x, z, dx, dz, n0, sigma_x, wavelength)

    # Plot
    fig, ax = plt.subplots(figsize=(8,6))
    im = ax.imshow(
        np.abs(E_out)**2,
        extent=[x[0], x[-1], z[0], z[-1]],
        origin='lower',
        aspect='auto',
        cmap='inferno'
    )
    ax.set_title("MMI Splitter BPM Propagation")
    ax.set_xlabel("x (µm)")
    ax.set_ylabel("z (µm)")
    fig.colorbar(im, ax=ax, label="Intensity")

    # Save data
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".npz")
    np.savez(tmp.name, E_out=E_out, x=x, z=z)
    tmp.close()

    return fig, tmp.name

def run_dispatcher(
    sim_type,
    # waveguide args
    w, l, L_bend,
    # mmi args
    z_start, L_MMI, w_MMI, w_wg, d,
    # common args
    n_WG, wavelength, ind_m
):
    if sim_type == "Waveguide S-bend":
        return run_waveguide(w, l, L_bend, n_WG, wavelength, ind_m)
    else:
        return run_mmi(z_start, L_MMI, w_MMI, w_wg, d, n_WG, wavelength, ind_m)

with gr.Blocks() as demo:
    gr.Markdown("## BPM Simulation Dashboard")

    # 1) Simulation selector
    sim_type = gr.Radio(
        choices=["Waveguide S-bend", "MMI Splitter"],
        value="Waveguide S-bend",
        label="Simulation Type"
    )

    with gr.Row():
        # 2) Waveguide parameters panel
        with gr.Column(visible=True) as waveguide_panel:
            w_slider     = gr.Slider(0.1, 5.0,   value=1.0,  step=0.1, label="Waveguide width w (µm)")
            l_slider     = gr.Slider(0.0, 10.0,  value=5.0,  step=0.1, label="Lateral offset l (µm)")
            Lb_slider    = gr.Slider(50.0, 500.0,value=200.0,step=10.0,label="S-bend length L (µm)")

        # 3) MMI parameters panel
        with gr.Column(visible=False) as mmi_panel:
            z0_slider    = gr.Slider(0.0, 200.0, value=50.0, step=1.0, label="MMI start z₀ (µm)")
            Lmmi_slider  = gr.Slider(10.0, 300.0,value=130.0,step=5.0, label="MMI length L (µm)")
            wmmi_slider  = gr.Slider(2.0, 20.0, value=8.0,  step=0.5, label="MMI width w_MMI (µm)")
            wwg_slider   = gr.Slider(0.5, 5.0,  value=2.0,  step=0.1, label="I/O waveguide width w_wg (µm)")
            d_slider     = gr.Slider(1.0, 20.0, value=4.0,  step=0.5, label="Waveguide separation d (µm)")

        # 4) Common parameters
        with gr.Column(scale=1):
            n_WG_slider      = gr.Slider(1.0, 2.0, value=1.1, step=0.01, label="Core refractive index n_WG")
            wavelength_slider= gr.Slider(0.4, 1.6, value=0.532,step=0.01, label="Wavelength λ (µm)")
            ind_m_slider     = gr.Slider(0,   4,   value=0,    step=1,    label="Mode index ind_m")

            run_button       = gr.Button("Run BPM")
            download_button  = gr.DownloadButton(label="Download data")

        # 5) Plot output
        with gr.Column(scale=2):
            plot_output = gr.Plot()

    # 6) Toggle panels on sim_type change
    sim_type.change(
        fn=lambda choice: (
            gr.update(visible=(choice=="Waveguide S-bend")),
            gr.update(visible=(choice=="MMI Splitter"))
        ),
        inputs=sim_type,
        outputs=[waveguide_panel, mmi_panel]
    )

    # 7) Wire up Run button + auto-update
    all_inputs = [
        sim_type,
        w_slider, l_slider, Lb_slider,
        z0_slider, Lmmi_slider, wmmi_slider, wwg_slider, d_slider,
        n_WG_slider, wavelength_slider, ind_m_slider
    ]
    run_button.click(
        fn=run_dispatcher,
        inputs=all_inputs,
        outputs=[plot_output, download_button]
    )
    for inp in all_inputs:
        inp.change(
            fn=run_dispatcher,
            inputs=all_inputs,
            outputs=[plot_output, download_button]
        )

demo.launch()