import logging
import warnings

import torchgmm.base

# This is taken from PyTorch Lightning and ensures that logging for this package is enabled
_root_logger = logging.getLogger()
_logger = logging.getLogger(__name__)
_logger.setLevel(logging.INFO)
if not _root_logger.hasHandlers():
    _logger.addHandler(logging.StreamHandler())
    _logger.propagate = False

# This disables most logs generated by PyTorch Lightning
logging.getLogger("lightning.pytorch").setLevel(logging.WARNING)
warnings.filterwarnings(action="ignore", message=".*Consider increasing the value of the `num_workers` argument.*")
warnings.filterwarnings(action="ignore", message=".*`LightningModule.configure_optimizers` returned `None`.*")
warnings.filterwarnings(action="ignore", message=".*`LoggerConnector.gpus_metrics` was deprecated in v1.5.*")


# We also want to define a function which silences info logs
def set_logging_level(level: int) -> None:
    """
    Enables or disables logging for the entire module. By default, logging is enabled.

    Args:
        enabled: Whether to enable logging.
    """
    _logger.setLevel(level)


# Export
__all__ = ["set_logging_level"]
