"""Tensorboard logger with add image interface."""

# Copyright (C) 2022-2024 Intel Corporation
# SPDX-License-Identifier: Apache-2.0


from pathlib import Path

import numpy as np
from matplotlib.figure import Figure

try:
    from lightning.pytorch.loggers.tensorboard import TensorBoardLogger
except ModuleNotFoundError:
    print("To use tensorboard logger install it using `pip install tensorboard`")
from lightning.pytorch.utilities import rank_zero_only

from .base import ImageLoggerBase


class AnomalibTensorBoardLogger(ImageLoggerBase, TensorBoardLogger):
    """Logger for tensorboard.

    Adds interface for `add_image` in the logger rather than calling the experiment object.

    .. note::
        Same as the Tensorboard Logger provided by PyTorch Lightning and the doc string is reproduced below.

    Logs are saved to
    ``os.path.join(save_dir, name, version)``. This is the default logger in Lightning, it comes
    preinstalled.

    Example:
        >>> from anomalib.engine import Engine
        >>> from anomalib.loggers import AnomalibTensorBoardLogger
        ...
        >>> logger = AnomalibTensorBoardLogger("tb_logs", name="my_model")
        >>> engine =  Engine(logger=logger)

    Args:
        save_dir (str): Save directory
        name (str | None): Experiment name. Defaults to ``'default'``.
            If it is the empty string then no per-experiment subdirectory is used.
            Default: ``'default'``.
        version (int | str | None): Experiment version. If version is not
            specified the logger inspects the save directory for existing
            versions, then automatically assigns the next available version.
            If it is a string then it is used as the run-specific subdirectory
            name, otherwise ``'version_${version}'`` is used.
            Defaults to ``None``
        log_graph (bool): Adds the computational graph to tensorboard. This
            requires that the user has defined the `self.example_input_array`
            attribute in their model.
            Defaults to ``False``.
        default_hp_metric (bool): Enables a placeholder metric with key
            ``hp_metric`` when ``log_hyperparams`` is called without a metric
            (otherwise calls to log_hyperparams without a metric are ignored).
            Defaults to ``True``.
        prefix (str): A string to put at the beginning of metric keys.
            Defaults to ``''``.
        **kwargs: Additional arguments like `comment`, `filename_suffix`, etc.
            used by :class:`SummaryWriter` can be passed as keyword arguments in
            this logger.
    """

    def __init__(
        self,
        save_dir: str,
        name: str | None = "default",
        version: int | str | None = None,
        log_graph: bool = False,
        default_hp_metric: bool = True,
        prefix: str = "",
        **kwargs,
    ) -> None:
        super().__init__(
            save_dir,
            name=name,
            version=version,
            log_graph=log_graph,
            default_hp_metric=default_hp_metric,
            prefix=prefix,
            **kwargs,
        )
        Path(save_dir).mkdir(parents=True, exist_ok=True)

    @rank_zero_only
    def add_image(self, image: np.ndarray | Figure, name: str | None = None, **kwargs) -> None:
        """Interface to add image to tensorboard logger.

        Args:
            image (np.ndarray | Figure): Image to log
            name (str | None): The tag of the image
                Defaults to ``None``.
            kwargs: Accepts only `global_step` (int). The step at which to log the image.
        """
        if "global_step" not in kwargs:
            msg = "`global_step` is required for tensorboard logger"
            raise ValueError(msg)

        # Need to call different functions of `SummaryWriter` for  Figure vs np.ndarray
        if isinstance(image, Figure):
            self.experiment.add_figure(figure=image, tag=name, close=False, **kwargs)
        else:
            self.experiment.add_image(img_tensor=image, tag=name, dataformats="HWC", **kwargs)
