import torch


def polynomial_fit(x, y, num_terms=5, weights=None):
    """
    Fits a polynomial curve of order "num_terms" to the datapoints in "x" and "y". This function
    supports channeling for fitting to mulitple independent data channels at once, and expects data
    in channeled form with shape like [channels, points]. (one channel is fine)
    Returns Tensor of shape [channels, num_terms].
    """
    X_cols = [x.pow(n) for n in range(num_terms)]
    X = torch.stack(X_cols, -1)

    if weights is None:
        a = torch.bmm(
            torch.linalg.solve(torch.bmm(X.transpose(1, 2), X), X.transpose(1, 2)),
            y.unsqueeze(-1),
        ).flatten(1)
    else:
        a = torch.bmm(
            torch.linalg.solve(
                torch.bmm(X.transpose(1, 2), torch.bmm(weights, X)), X.transpose(1, 2)
            ),
            torch.bmm(weights, y.unsqueeze(-1)),
        ).flatten(1)
    return a