from pathlib import Path
import shutil
from typing import Any, Optional

import numpy as np
import torch
from PIL import Image
from pytorch_lightning.loggers.logger import Logger
from pytorch_lightning.utilities import rank_zero_only


class LocalLogger(Logger):
    def __init__(self, log_path: Path | str = Path("outputs/local")) -> None:
        super().__init__()
        self.experiment = None
        self.log_path = Path(log_path)
        shutil.rmtree(self.log_path, ignore_errors=True)

    @property
    def name(self):
        return "LocalLogger"

    @property
    def version(self):
        return 0

    @rank_zero_only
    def log_hyperparams(self, params):
        pass

    @rank_zero_only
    def log_metrics(self, metrics, step):
        pass

    @rank_zero_only
    def log_image(
        self,
        key: str,
        images: list[Any],
        step: Optional[int] = None,
        **kwargs,
    ):
        # The function signature is the same as the wandb logger's, but the step is
        # actually required.
        assert step is not None
        for index, image in enumerate(images):
            path = self.log_path / f"{key}/{index:0>2}_{step:0>6}.png"
            path.parent.mkdir(exist_ok=True, parents=True)
            if isinstance(image, torch.Tensor):
                Image.fromarray(image.permute(1, 2, 0).numpy().astype(np.uint8)).save(path)
            else:
                Image.fromarray(image).save(path)