import matplotlib.pyplot as plt
from .HepFigure import HepAxes

__all__ = ["HepAxes", "enable_default_hep_subplots"]


_ORIGINAL_SUBPLOTS = plt.subplots
_HEP_DEFAULT_ENABLED = False


def enable_default_hep_subplots(enable=True):
	"""Enable or disable hep projection as the default for ``plt.subplots``.

	When enabled, ``plt.subplots()`` will behave like:
	``plt.subplots(subplot_kw={"projection": "hep"})`` unless the caller
	explicitly provides a projection in ``subplot_kw``.
	"""

	global _HEP_DEFAULT_ENABLED

	if enable and not _HEP_DEFAULT_ENABLED:
		def _hep_default_subplots(*args, **kwargs):
			subplot_kw = dict(kwargs.get("subplot_kw") or {})
			subplot_kw.setdefault("projection", "hep")
			kwargs["subplot_kw"] = subplot_kw
			return _ORIGINAL_SUBPLOTS(*args, **kwargs)

		plt.subplots = _hep_default_subplots
		_HEP_DEFAULT_ENABLED = True
	elif not enable and _HEP_DEFAULT_ENABLED:
		plt.subplots = _ORIGINAL_SUBPLOTS
		_HEP_DEFAULT_ENABLED = False


# Default-on so users can call plt.subplots() directly after importing HEP_PLOT.
enable_default_hep_subplots(True)
