from pyscript import display
import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import curve_fit


def func_linear(x, m, b):
    return m * x + b


def scatter_plot(x, y, title, xlabel="Voltage", ylabel="Current"):
    fig, ax = plt.subplots()
    ax.scatter(-1 * x, y)
    ax.set_title(title)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.grid(True, alpha=0.25)
    display(fig, target="live-figures", append=True)
    plt.close(fig)


# Current-voltage data sets from the notebook.
scatter_plot(
    np.array([-3, -2.5, -2, -1.5, -1, -0.5, 0, 0.5, 1, 1.5, 2, 2.5, 3, -0.1, -0.2, -0.3, -0.4]),
    np.array([-1.52, -1.49, -1.49, -1.47, -1.46, -1.26, 7.1, 19.4, 28.1, 37.4, 45.0, 50.6, 54.0, 4.13, 1.87, 0.06, -0.89]),
    "Trial 1 current-voltage data",
)

scatter_plot(
    np.array([-2, -1.9, -1.8, -1.7, -1.6, -1.5, -1.4, -1, 0, 0.5]),
    np.array([-1.5, -1.36, -1.04, -0.3, 1.26, 3.85, 8.0, 31.4, 123, 170]),
    "Trial 2 current-voltage data",
)

x = np.array([-2, -1.9, -1.8, -1.7, -1.6, -1.5, -1.4, -1.3, -1, 0, 1])
y = np.array([-1.33, -1.30, -1.26, -1.22, -1.12, -0.94, -0.53, 0.25, 5.8, 55, 112])
popt, pcov = curve_fit(func_linear, x, y)
scatter_plot(x, y, "Trial 3 current-voltage data")

scatter_plot(
    np.array([-2, -1.5, -1.4, -1.3, -1.2, -1.1, -1, -0.9, -0.5, 0]),
    np.array([-1.53, -1.44, -1.38, -1.24, -0.85, -0.05, 1.34, 3.7, 17.9, 45]),
    "Trial 4 current-voltage data",
)

scatter_plot(
    np.array([-2, -1.5, -1, -0.9, -0.8, -0.7, -0.6, -0.5, -0.4, -0.3, 0, 1]),
    np.array([-1.8, -1.79, -1.77, -1.76, -1.72, -1.66, -1.41, -0.78, 0.41, 1.92, 10.6, 43]),
    "Trial 5 current-voltage data",
)

# Stopping potential fit from the notebook.
x = np.array([5.2, 8.22, 7.41, 6.88, 5.49])
y = np.array([0.4, 1.7, 1.4, 1.3, 0.7])
yerr = np.array([0.02, 0.01, 0.01, 0.005, 0.01])
perr = np.sqrt(np.diag(pcov))
m_error = perr[0]

fit, fit_cov = curve_fit(func_linear, x, y, sigma=yerr, absolute_sigma=True)
fig, ax = plt.subplots()
ax.errorbar(x, y, yerr, fmt="o", linewidth=2, capsize=6)
ax.plot(x, func_linear(x, *fit))
ax.set_title("Stopping potential fit")
ax.set_xlabel("Frequency")
ax.set_ylabel("Stopping potential")
ax.grid(True, alpha=0.25)
display(fig, target="live-figures", append=True)
plt.close(fig)

display(f"m={popt[0]}\nSlope Uncertainty: {m_error}", target="live-output")