import torch


def _parabolic_trend_crosscov(S1, L1, S2, L2, gamma0, sigma0s, betas, s0):
    """ Same as above, but for non-stationary variance models.

    The model is the following:
        \sigma_i(s) = \sigma0_i + \beta_i * (s - s0)**2

    The genial simplicity of the implementation is this: instead of having an
    array containing the variances of each response, have an array of functions
    giving the variances of each function at some place given as argument to
    the function.
    The beauty of having functions as first-class citizens in Python has rarely
    shone so brightly.

    Parameters
    ----------
    S1: (M, d) Tensor
        Spatial location vector. Note if d=1, should still have two
        dimensions.
    L1: (M) Tensor
        Response indices vector.
    S2: (N, d) Tensor
        Spatial location vector.
    L2: (N) Tensor
        Response indices vector.
    gamma0: Tensor
        Coupling parameter.
    sigma0s: (p) Tensor
        Vector of individual standard deviation (constant part) for each to the p components.
    betas: (p, n_dim, n_dim) Tensor
        Parabolic trend defined by a quadratic form for each component.
        The final model is
            sigma_i(s) = sigma0s[i] + (s - s0)**2 * betas[i, :, :] * (s - s0)**2.
    s0: (n_dim) Tensor
        Reference point in the domain for linear trend.

    Returns
    -------
    gamma: (M, N) Tensor

    """
    # Turn to matrices of size (M, N).
    L1mat, L2mat = torch.meshgrid(L1, L2)

    # Same for the spatiality. Matrices of size (M, N, n_dim).
    S1mat, _ = torch.meshgrid(S1.reshape(-1), S2[:, 0])
    _, S2mat = torch.meshgrid(S1[:, 0], S2.reshape(-1))

    S1mat = S1mat.reshape(S1.shape[0], S1.shape[1], S2.shape[0]).transpose(1,2)
    S2mat = S2mat.reshape(S1.shape[0], S2.shape[0], S2.shape[1])

    # Coupling part.
    # Have to extract the float value from gamma0 to use fill.
    gamma_mat = (torch.Tensor(L1mat.shape).fill_(gamma0.item())
            + (1- gamma0) * (L1mat == L2mat))

    # Notice the GENIUS of Pytorch: If we want A_ij to contain sigma[Aij]
    # we just do simga[A] and get the whole matrix, with the same shape as A.
    # This is beautiful.
    sigma0_mat1, sigma0_mat2 = sigma0s[L1mat], sigma0s[L2mat]

    # Fetch the spatial coefficient vectors.
    beta_mat1, beta_mat2 = betas[L1mat], betas[L2mat]

    # Compute the quadratic form.
    diff1 = S1mat - s0
    diff2 = S2mat - s0
    dot_mat1 = torch.einsum("ija,ijab,ijb->ij", diff1, beta_mat1, diff1)
    dot_mat2 = torch.einsum("ija,ijab,ijb->ij", diff2, beta_mat2, diff2)

    return (sigma0_mat1 + dot_mat1) * (sigma0_mat2 + dot_mat2) * gamma_mat