import torch


def to_monochrome(x: torch.Tensor,
                  color: str,
                  preserve_energy:bool=False) -> torch.Tensor:
    """
    Transform a single-channel grayscale (3dim, (1,h,w) tensor)
    to a mono-chrome 3channel tensor (3,h,w), of either gray, red, green, blue.
    - If color is one of [red, green,blue], then outputs a 3-channel tensor by
    putting the input single-channel into the proper channel.
    - If color is gray, then distributes the input grayscale tensor equally into
    the r,g,b channels (ie. puts the input/3 into each r,g,b channels)

    Args:
    x: a single 3dim torch.Tensor a single channel; (1, h, w)
    color: str - one of ['red', 'green', 'blue']

    returns:
    - a single (3dim) torch.Tensor with 3 channels: (3, h, w)

    """
    out = torch.zeros_like(x).repeat((3,1,1))

    color = color.lower()
    color2dim = {"red": 0, "green": 1, "blue": 2}
    if color == 'gray':
        factor = 3.0 if preserve_energy else 1.0
        for i in range(3):
            out[i] = x/factor
    else:
        color_dim = color2dim[color]
        out[color_dim] = x

    return out