import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
from scipy import special as sp


def plot_function(function_name, x_min, x_max, y_min, y_max):
    functions = {
        "abs": np.abs,
        "acos": np.arccos,
        "acosh": np.arccosh,
        "airy": sp.airy,
        "asin": np.arcsin,
        "asinh": np.arcsinh,
        "atan": np.arctan,
        "atanh": np.arctanh,
        "ceil": np.ceil,
        "cos": np.cos,
        "cosh": np.cosh,
        "erf": lambda x: np.vectorize(sp.erfi)(x).real,
        "erfc": lambda x: np.vectorize(sp.erfi)(x).imag,
        "exp": np.exp,
        "floor": np.floor,
        "imag": np.imag,
        "int": lambda x: np.vectorize(int)(x),
        "inverf": lambda x: np.vectorize(sp.erfi)(x).real,
        "invnor": lambda x: np.vectorize(sp.erfi)(x).imag,
        "lambertw": lambda x: np.vectorize(sp.lambertw)(x).real,
        "log": np.log,
        "log10": np.log10,
        "real": np.real,
        "sgn": lambda x: np.sign(x),
        "sin": np.sin,
        "sinh": np.sinh,
        "sqrt": np.sqrt,
        "tanh": np.tanh,
        "jv": sp.jv,
        "EllipticK": sp.ellipk,
        "EllipticE": sp.ellipe,
        "erfi": sp.erfi,
        "gamma": sp.gamma,
        "tan": np.tan
    }
    x = np.linspace(x_min, x_max, 1000)
    y = functions[function_name](x)
    plt.ylim(y_min, y_max)
    plt.plot(x, y)
    plt.title(function_name)
    plt.xlabel('x')
    plt.ylabel('y')
    plt.grid(True)
    plt.tight_layout()

    return plt.gcf()

evals = ["abs", "acos", "acosh", "airy", "asin", "asinh", "atan", "atanh", "cdawson", "ceil", "cerf", "cos", "cosh", "erf", "erfc", "exp", "floor", "imag", "int", "inverf", "invnor", "lambertw", "lgamma", "log", "log10", "real", "sgn", "sin", "sinh", "sqrt", "tanh", "jv", "besj0", "besj1", "besy0", "EllipticK", "EllipticE", "besy1", "erfi", "gamma", "tan"]

inputs = [
    gr.Dropdown(choices=evals),
    gr.Number(0, label="x min"),
    gr.Number(1, label="x max"),
    gr.Number(0, label="y min"),
    gr.Number(1, label="y max"),
]
iface = gr.Interface(
    title='gnuplot',
    fn=plot_function,
    inputs=inputs,
    outputs="plot",
    live=True
)
iface.launch()