import torch


def f_get_raster_image(cfg,
                       images,
                       history_weight=0.9):
    """
    Creates single raster image from sequence of images from l5kit's AgentDataset
        Args:
            cfg {dict}: Dictionary config.
            images: (batch_size, 2*(history_num_frames+1)+3, H, H) - sequences of images after applying l5kit's rasterizer:
                    There is (history_num_frames+1) ego-car images, (history_num_frames+1) agent-car's images + 3 scene RGB images
            history_weight {float}: Amount of history fading (for rendering).

        Returns:
            RGB image of the scene and agents.
            Red color stays for EGO car.
            Yellow color stays for Agent's cars.
    """

    batch_size = images.shape[0]
    image_size = images.shape[-1]

    # get number of history steps
    hnf = cfg['model_params']['history_num_frames']

    # define ego-car's indices range in channels (images):
    # ind (0, hnf) correspond to all agents except ego car,
    # from hnf+1 to 2*hnf+1 correspond to ego car,
    # last 3 indices correspond to rgb scene
    ego_index = range(hnf+1, 2*hnf+2)

    # iterate through ego-car's frames and sum them according to history_weight (history fading) in single channel.
    ego_path_image = torch.zeros(size=(batch_size, image_size, image_size), device=cfg['device'])
    for im_id in reversed(ego_index):
        ego_path_image = (images[:, im_id, :, :] + ego_path_image * history_weight).clamp(0, 1)

    # define agent's range
    agents_index = range(0, hnf+1)

    # iterate through agent-car's frames and sum them according to history_weight in single channel
    agents_path_image = torch.zeros(size=(batch_size, image_size, image_size), device=cfg['device'])
    for im_id in reversed(agents_index):
        agents_path_image = (images[:, im_id, :, :] + agents_path_image*history_weight).clamp(0, 1)

    #  RGB path for ego (red (255, 0, 0)); channels last
    ego_path_image_rgb = torch.zeros((ego_path_image.shape[0],
                                      ego_path_image.shape[1],
                                      ego_path_image.shape[2],
                                      3), device=cfg['device'])

    ego_path_image_rgb[:, :, :, 0] = ego_path_image

    # RGB paths for agents (yellow (255, 255, 0)); channels last
    agents_path_image_rgb = torch.zeros((agents_path_image.shape[0],
                                         agents_path_image.shape[1],
                                         agents_path_image.shape[2],
                                         3), device=cfg['device'])
    # yellow
    agents_path_image_rgb[:, :, :, 0] = agents_path_image
    agents_path_image_rgb[:, :, :, 1] = agents_path_image

    # generate full RGB image with all cars (ego + agents)
    all_vehicles_image = ego_path_image_rgb + agents_path_image_rgb  # (batch_size, 3, H, H)

    # get RGB image for scene from rasterizer (3 last images); channels last
    scene_image_rgb = images[:, 2*hnf+2:, :, :].permute(0, 2, 3, 1)

    # Add mask to positions of cars (to merge all layers):
    # We need to take into account that the off-road is white, i.e. 1 in all channels
    # So, we have to prevent disappearing of off-road cars after clipping when we add images together.
    # For ex. (1, 1, 1) + (1, 0, 0) = (2, 1, 1) --> (1, 1, 1) = off-road car disappears.
    # In order to solve this, we cut the scene at the car's area.

    scene_image_rgb[(all_vehicles_image > 0).any(dim=-1)] = 0.0

    # generate final raster map
    full_raster_image = (all_vehicles_image + scene_image_rgb).clamp(0, 1)

    # channels as a second dimension
    full_raster_image = full_raster_image.permute(0, 3, 1, 2)

    return full_raster_image