import numpy as np
from scipy import optimize
import matplotlib.pyplot as plt

Te_ref = np.loadtxt("temperature.csv", delimiter=",")
Te_ref[:, 0] *= 1e-4  # Convert X axis to m


def te_function(ypos, mid, wwid, w0, w1, w2, Tmax, Tmin, clip=False):
    width = w0 + ((ypos - mid) * w1 + (ypos - mid) ** 2 * w2) * np.exp(
        -(((ypos - mid) / wwid) ** 2)
    )

    if clip:
        width = np.clip(width, 1e-10, None)

    return Tmax - 0.5 * (1 + np.tanh((ypos - mid) / width)) * (Tmax - Tmin)


popt, pcov = optimize.curve_fit(
    te_function,
    Te_ref[:, 0],
    Te_ref[:, 1],
    p0=[2.2e-4, 1e-4, 1e-4, 0.0, 0.0, 0.960, 0.190],
)

print(popt)

xfit = np.linspace(Te_ref[0, 0], Te_ref[-1, 0], 100)

plt.plot(xfit, te_function(xfit, *popt, clip=True), "-k")
plt.plot(Te_ref[:, 0], Te_ref[:, 1], "or")
plt.show()
