Logger

def clear_logs(self):
    """
    Clear logs.

    Examples:
        >>> logger = Logger()
        >>> logger.clear_logs()
    """
    self.logs = []
    if self.verbose:
        print("Logs cleared.")

def clear_snapshots(self):
    """
    Clear snapshots.

    Examples:
        >>> logger = Logger()
        >>> logger.clear_snapshots()
    """
    self.snapshots = []
    if self.verbose:
        print("Snapshots cleared.")

def load_snapshots(self, snapshots_filepath: str):
    """
    Load snapshots from disk.

    Parameters:
        snapshots_filepath:
            Filepath for snapshots in `.pth` format.

    Raises:
        FileNotFoundError:
            If snapshots file not found.

    Examples:
        >>> logger = Logger()
        >>> logger.load_snapshots()
    """
    if not os.path.exists(snapshots_filepath):
        raise FileNotFoundError(f"Snapshots file not found: {snapshots_filepath}")

    self.snapshots = torch.load(snapshots_filepath)
    if self.verbose:
        print(f"Snapshots loaded: {snapshots_filepath}")

def log_image_to_tensorboard(self, **kwargs: Any):
    """
    Log image to tensorboard.

    Parameters:
        **kwargs: Scalars to log to tensorboard.

    Shape:
        img_tensor: Default is :math:`(3, H, W)`. You can use ``torchvision.utils.make_grid()`` to
        convert a batch of tensor into 3xHxW format or call ``add_images`` and let us do the job.
        Tensor with :math:`(1, H, W)`, :math:`(H, W)`, :math:`(H, W, 3)` is also suitable as long as
        corresponding ``dataformats`` argument is passed, e.g. ``CHW``, ``HWC``, ``HW``.

    Examples:
        >>> logger = Logger()
        >>> logger.log_image_to_tensorboard(tag='state',
                img_tensor=torch.rand(3, 128, 128),
                global_step=0)
    """
    if self.tensorboard_writer is not None:
        self.tensorboard_writer.add_image(**kwargs)

def log_images_to_tensorboard(self, **kwargs: Any):
    """
    Log images to tensorboard.

    Parameters:
        **kwargs: Scalars to log to tensorboard.

    Shape:
        img_tensor: Default is :math:`(N, 3, H, W)`. If ``dataformats`` is specified, other shape will be
        accepted. e.g. NCHW or NHWC.

    Examples:
        >>> logger = Logger()
        >>> logger.log_images_to_tensorboard(tag='state',
                img_tensor=torch.rand(10, 3, 128, 128),
                global_step=0)
    """
    if self.tensorboard_writer is not None:
        self.tensorboard_writer.add_images(**kwargs)

def log_p_burn(self, simulator: 'WildTorchSimulator'):
    """
    Log p(burn) of the simulation.

    Parameters:
        simulator: Wildfire simulator instance.

    Examples:
        >>> logger = Logger()
        >>> logger.log_p_burn(simulator)
    """
    self.p_burns.append(simulator.p_burn.detach().clone())

def log_scalars_to_tensorboard(self, step: int = None, **kwargs: Any):
    """
    Log scalars to tensorboard.

    Parameters:
        step: Simulation step.
        **kwargs: Scalars to log to tensorboard.

    Examples:
        >>> logger = Logger()
        >>> logger.log_scalars_to_tensorboard(step=0, burning_cells={
                'max': 100,
                'min': 0,
                'mean': 50,
            })
    """
    if self.tensorboard_writer is not None:
        for key, value in kwargs.items():
            self.tensorboard_writer.add_scalars(key, value, step)

def log_stats(self, step: int = None, **kwargs: Any):
    """
    Log simulation statistics.

    Parameters:
        step: Simulation step.
        **kwargs: Simulation statistics to log.

    Examples:
        >>> logger = Logger()
        >>> logger.log_stats(step=0, burning_cells=100, burned_cells=50)
    """
    log_entry = {key: value for key, value in kwargs.items()}
    if step is not None:
        log_entry["step"] = step
    self.logs.append(log_entry)

    if self.verbose:
        print(log_entry)

    # Tensorboard
    if self.tensorboard_writer is not None:
        for key, value in kwargs.items():
            self.tensorboard_writer.add_scalar(key, value, step)

def log_video_to_tensorboard(self, **kwargs: Any):
    """
    Log video to tensorboard.

    Needs `moviepy`.

    Parameters:
        **kwargs: Scalars to log to tensorboard.

    Shape:
        vid_tensor: :math:`(N, T, C, H, W)`.
        The values should lie in [0, 255] for type `uint8` or [0, 1] for type `float`.

    Examples:
        >>> logger = Logger()
        >>> logger.log_video_to_tensorboard(tag='state',
                vid_tensor=torch.rand(1, 10, 3, 128, 128),
                global_step=0)
    """
    if self.tensorboard_writer is not None:
        self.tensorboard_writer.add_video(**kwargs)

def save_logs(self):
    """
    Save logs to disk.

    Examples:
        >>> logger = Logger()
        >>> logger.save_logs()
    """
    if self.disable_writing:
        return

    if self.verbose:
        print(f"Logs saved: {self.json_filepath}")
    if self.tensorboard_writer is not None:
        self.tensorboard_writer.flush()
        if self.verbose:
            print(f"Tensorboard logs saved: {self.tensorboard_writer.log_dir}")

def save_snapshots(self):
    """
    Save snapshots to disk.

    Examples:
        >>> logger = Logger()
        >>> logger.save_snapshots()
    """
    torch.save(self.snapshots, self.snapshots_filepath)
    if self.verbose:
        print(f"Snapshots saved: {self.snapshots_filepath}")

def snapshot_simulation(self, simulator: 'WildTorchSimulator'):
    """
    Take a snapshot of the simulation.

    Parameters:
        simulator: Wildfire simulator instance.

    Examples:
        >>> logger = Logger()
        >>> logger.snapshot_simulation(simulator)
    """
    self.snapshots.append(simulator.checkpoint)