import torch


def linear_2_oklab(x):
    """Converts pytorch tensor 'x' from Linear to OkLAB colorspace, described here:
        https://bottosson.github.io/posts/oklab/
    Inputs:
        x -- pytorch tensor of size B x 3 x H x W, assumed to be in linear 
             srgb colorspace, scaled between 0. and 1.
    Returns:
        y -- pytorch tensor of size B x 3 x H x W in OkLAB colorspace
    """
    assert x.size(1) == 3, "attempted to convert colorspace of tensor w/ > 3 channels"

    x = torch.clamp(x, 0., 1.)

    r = x[:, 0:1, :, :]
    g = x[:, 1:2, :, :]
    b = x[:, 2:3, :, :]

    li = 0.4121656120 * r + 0.5362752080 * g + 0.0514575653 * b
    m = 0.2118591070 * r + 0.6807189584 * g + 0.1074065790 * b
    s = 0.0883097947 * r + 0.2818474174 * g + 0.6302613616 * b

    li = torch.pow(li, 1. / 3.)
    m = torch.pow(m, 1. / 3.)
    s = torch.pow(s, 1. / 3.)

    L = 0.2104542553 * li + 0.7936177850 * m - 0.0040720468 * s
    A = 1.9779984951 * li - 2.4285922050 * m + 0.4505937099 * s
    B = 0.0259040371 * li + 0.7827717662 * m - 0.8086757660 * s

    y = torch.cat([L, A, B], 1)
    return y