import torch


def generate_novel_views(model, img_source, azimuth_source, elevation_source,
                         azimuth_shifts, elevation_shifts):
    """Generates novel views of an image by inferring its scene representation,
    rotating it and rendering novel views. Returns a batch of images
    corresponding to the novel views.

    Args:
        model (models.neural_renderer.NeuralRenderer): Neural rendering model.
        img_source (torch.Tensor): Single image. Shape (channels, height, width).
        azimuth_source (torch.Tensor): Azimuth of source image. Shape (1,).
        elevation_source (torch.Tensor): Elevation of source image. Shape (1,).
        azimuth_shifts (torch.Tensor): Batch of angle shifts at which to
            generate novel views. Shape (num_views,).
        elevation_shifts (torch.Tensor): Batch of angle shifts at which to
            generate novel views. Shape (num_views,).
    """
    # No need to calculate gradients
    with torch.no_grad():
        num_views = len(azimuth_shifts)
        # Batchify image
        img_batch = img_source.unsqueeze(0)
        # Infer scene
        scenes = model.inverse_render(img_batch)
        # Copy scene for each target view
        scenes_batch = scenes.repeat(num_views, 1, 1, 1, 1)
        # Batchify azimuth and elevation source
        azimuth_source_batch = azimuth_source.repeat(num_views)
        elevation_source_batch = elevation_source.repeat(num_views)
        # Calculate azimuth and elevation targets
        azimuth_target = azimuth_source_batch + azimuth_shifts
        elevation_target = elevation_source_batch + elevation_shifts
        # Rotate scenes
        rotated = model.rotate_source_to_target(scenes_batch, azimuth_source_batch,
                                                elevation_source_batch,
                                                azimuth_target, elevation_target)
    # Render images
    return model.render(rotated).detach()