import numpy as np
from scipy.optimize import least_squares

def k_fit(y_dat):
    """
    Fit linear and quadratic to K total
    """
    def optimize_func(a, y, n):
        result = []
        i, j = 0, 0
        for u in range(len(y)):
            val = y[0] + a[0] * i + a[1] * j**2 - y[u]
            i = i + 1
            if u >= n: j = j + 1
            result.append(val)
        return result
    mini = None
    for n in range(len(y_dat)):
        a_init = (0.0001, 0.0001)
        result = least_squares(optimize_func, a_init, args=(y_dat, n))
        if mini is None or result["cost"] < mini:
            mini = result["cost"]
            n_fit = n
            a_fit = result["x"]
    return (a_fit[0], a_fit[1], n_fit)