"""Mixin classes for unittests"""
from chromatictools import math
import functools
import contextlib
import numpy as np
import io


class AssertPrintsMixin:
  """Mixin class for print assertion"""
  @contextlib.contextmanager
  def assert_prints(self, target: str):
    """Assert that contextual code prints the target string.
    Use as context manager

    Args:
      target (str): Print expectation"""
    with io.StringIO() as buf, contextlib.redirect_stdout(buf):
      yield
      printed = buf.getvalue()
    self.assertEqual(printed, target)


class AssertDoesntRaiseMixin:
  """Mixin class for "doesn't-raise" assertion"""
  @contextlib.contextmanager
  def assert_doesnt_raise(self):
    """Assert that contextual code doesn't raise any exception
    Use as context manager"""
    try:
      yield
    except Exception as e:
      raise AssertionError from e


class SignificantPlacesAssertMixin:
  """Mixin class for assertions based on singificant places"""
  def assert_almost_equal_significant(
    self,
    first,
    second,
    places=0,
    msg=None,
    delta=None
  ):
    """Fail if the two objects are unequal as determined by their
    difference rounded to the given number of significant decimal places
    (default 0)"""
    data = [first, second]
    if any(data):
      data = math.mantissa(data)
    return self.assertAlmostEqual(
      *data,
      places=places,
      msg=msg,
      delta=delta
    )


class RMSEAssertMixin:
  """Mixin class for assertions based on the RMSE"""
  def _assert_rmse(
    self,
    x: np.ndarray,
    y: np.ndarray,
    almost: bool,
    *args,
    rmse: float = 0,
    **kwargs
  ):
    """Assert that the RMSE is equal to a given value

    Arguments:
      x (array): First array
      y (array): Second array
      almost (bool): If True, then assert that the RMSE is almost zero
      rmse (float): RMSE value. Defaults to 0
      *args: Positional arguments for :func:`unittest.TestCase.assertEqual`
        or :func:`unittest.TestCase.assertAlmostEqual`
      **kwargs: Keyword arguments for :func:`unittest.TestCase.assertEqual`
        or :func:`unittest.TestCase.assertAlmostEqual`"""
    return (self.assertAlmostEqual if almost else self.assertEqual)(
      math.rmse(x, y), rmse,
      *args, **kwargs,
    )

  assert_equal_rmse = functools.partialmethod(_assert_rmse, almost=False)
  assert_almost_equal_rmse = functools.partialmethod(_assert_rmse, almost=True)

"""Mixin classes for unittests"""
from chromatictools import math
import functools
import contextlib
import numpy as np
import io


class AssertPrintsMixin:
  """Mixin class for print assertion"""
  @contextlib.contextmanager
  def assert_prints(self, target: str):
    """Assert that contextual code prints the target string.
    Use as context manager

    Args:
      target (str): Print expectation"""
    with io.StringIO() as buf, contextlib.redirect_stdout(buf):
      yield
      printed = buf.getvalue()
    self.assertEqual(printed, target)


class AssertDoesntRaiseMixin:
  """Mixin class for "doesn't-raise" assertion"""
  @contextlib.contextmanager
  def assert_doesnt_raise(self):
    """Assert that contextual code doesn't raise any exception
    Use as context manager"""
    try:
      yield
    except Exception as e:
      raise AssertionError from e


class SignificantPlacesAssertMixin:
  """Mixin class for assertions based on singificant places"""
  def assert_almost_equal_significant(
    self,
    first,
    second,
    places=0,
    msg=None,
    delta=None
  ):
    """Fail if the two objects are unequal as determined by their
    difference rounded to the given number of significant decimal places
    (default 0)"""
    data = [first, second]
    if any(data):
      data = math.mantissa(data)
    return self.assertAlmostEqual(
      *data,
      places=places,
      msg=msg,
      delta=delta
    )


class RMSEAssertMixin:
  """Mixin class for assertions based on the RMSE"""
  def _assert_rmse(
    self,
    x: np.ndarray,
    y: np.ndarray,
    almost: bool,
    *args,
    rmse: float = 0,
    **kwargs
  ):
    """Assert that the RMSE is equal to a given value

    Arguments:
      x (array): First array
      y (array): Second array
      almost (bool): If True, then assert that the RMSE is almost zero
      rmse (float): RMSE value. Defaults to 0
      *args: Positional arguments for :func:`unittest.TestCase.assertEqual`
        or :func:`unittest.TestCase.assertAlmostEqual`
      **kwargs: Keyword arguments for :func:`unittest.TestCase.assertEqual`
        or :func:`unittest.TestCase.assertAlmostEqual`"""
    return (self.assertAlmostEqual if almost else self.assertEqual)(
      math.rmse(x, y), rmse,
      *args, **kwargs,
    )

  assert_equal_rmse = functools.partialmethod(_assert_rmse, almost=False)
  assert_almost_equal_rmse = functools.partialmethod(_assert_rmse, almost=True)

Mixin classes for unittests Mixin class for print assertion Assert that contextual code prints the target string.     Use as context manager      Args:       target (str): Print expectation Mixin class for "doesn't-raise" assertion Assert that contextual code doesn't raise any exception     Use as context manager Mixin class for assertions based on singificant places Fail if the two objects are unequal as determined by their     difference rounded to the given number of significant decimal places     (default 0) Mixin class for assertions based on the RMSE Assert that the RMSE is equal to a given value      Arguments:       x (array): First array       y (array): Second array       almost (bool): If True, then assert that the RMSE is almost zero       rmse (float): RMSE value. Defaults to 0       *args: Positional arguments for :func:`unittest.TestCase.assertEqual`         or :func:`unittest.TestCase.assertAlmostEqual`       **kwargs: Keyword arguments for :func:`unittest.TestCase.assertEqual`         or :func:`unittest.TestCase.assertAlmostEqual`