# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import logging
import typing as tp

import dora
import torch


logger = logging.getLogger(__name__)


class Profiler:
    """Context manager wrapper for xformers profiler.
    """
    def __init__(self, module: torch.nn.Module, enabled: bool = False):
        self.profiler: tp.Optional[tp.Any] = None
        if enabled:
            from xformers.profiler import profile
            output_dir = dora.get_xp().folder / 'profiler_data'
            logger.info("Profiling activated, results with be saved to %s", output_dir)
            self.profiler = profile(output_dir=output_dir, module=module)

    def step(self):
        if self.profiler is not None:
            self.profiler.step()  # type: ignore

    def __enter__(self):
        if self.profiler is not None:
            return self.profiler.__enter__()  # type: ignore

    def __exit__(self, exc_type, exc_value, exc_tb):
        if self.profiler is not None:
            return self.profiler.__exit__(exc_type, exc_value, exc_tb)  # type: ignore