text: import os
  import logging
  import pickle
  import pytest
  import numpy as np
  from unittest import TestCase, util
  
  from astropy.coordinates import Angle
  import astropy.units as u
  
  from src.identify.richness import NFWHelper, NFWProjection, integrand_radial, \
      CorrectionRadial, SchecterFunction, CorrectionCompleteness, richness, \
      RadialFilter
  from src.models.make_models import load_training_data
  from src.objects import load_data
  from src.utils import Str
  import src.globals as glo
  
  
  def load_test_candidate(test_guid, dataset):
      """
      From pickled CandidateList Object
  
      Parameters
      ----------
      test_guid
      dataset
  
      Returns
      -------
  
      """
      name = '{0}_CandidateList.pkl'.format(dataset)
      data = os.path.join('data', 'interim', name)
      with open(data, "rb") as data_in:
          return pickle.load(data_in)[test_guid]
  
  
  def load_test_candidate_gold(test_guid):
      """
      From pickled Candidate Object
  
      Parameters
      ----------
      test_guid
  
      Returns
      -------
  
      """
  
      name_gold = '{0}_Candidate.pkl'.format(test_guid)
      data_gold = os.path.join('data', 'test_candidates', 'pickle', name_gold)
      with open(data_gold, "rb") as data_in_gold:
          return pickle.load(data_in_gold)
  
  
  class TestNFWHelper(TestCase):
      def setUp(self):
          """
          Tests an np.array covering the following scenarios:
          x[0] > 1, x[1] < 0, x[2] == 1, x[3] < 1, x[4] == 0
          """
          self.ph_np = NFWHelper(x=np.array([1.5, -0.2, 1, 0.3, 0]))
  
      def test_denominator(self):
          expectation = np.array(
              [1.11803399, np.nan, np.nan, 0.953939201, np.nan])
          assert self.ph_np.denominator == pytest.approx(expectation, nan_ok=True)
  
      def test_help(self):
          expectation = np.array(
              [0.24772531154589, np.nan, np.nan, -0.9642973469840, np.nan])
          assert self.ph_np.help == pytest.approx(expectation, nan_ok=True)
  
  
  class TestProjectedNFW(TestCase):
      def setUp(self):
          self.nfw1 = NFWProjection(radius=2.5, r500=0.3)
          self.nfw2 = NFWProjection(radius=2.5, scale_radius=0.2)
          self.nfw3 = NFWProjection(radius=15, scale_radius=10)
          self.nfw4 = NFWProjection(radius=3, scale_radius=10)
  
      def test_scale_radius_radius_greater_than_r500_no_scale_radius(self):
          assert self.nfw1.scale_radius == pytest.approx(0.1, rel=1e-8)
  
      def test_normalised_radius_scale_radius2(self):
          assert self.nfw2.x == pytest.approx(12.5, rel=1e-5)
  
      def test_normalised_radius_scale_radius4(self):
          assert self.nfw4.x == pytest.approx(0.3, rel=1e-5)
  
      def test_denominator_radius_greater_than_scaled_radius_no_r500(self):
          assert self.nfw2.denominator == pytest.approx(155.25, rel=1e-5)
  
      def test_no_scaled_radius_no_r500(self):
          with pytest.raises(ValueError):
              NFWProjection(radius=2.5)
  
      def test_profile_radius_gt_scale_radius(self):
          assert self.nfw3.profile == pytest.approx(0.198180249, rel=1e-5)
  
      def test_profile_radius_lt_scale_radius(self):
          assert self.nfw4.profile == pytest.approx(1.05966741, rel=1e-5)
  
  
  class TestIntegrandRadial(TestCase):
      def setUp(self):
          """
          Tests an np.array covering the following scenarios:
          x = R / R_s
          x[0] > 1, x[1] < 0, x[2] == 1, x[3] < 1, x[4] == 0
          """
          self.r_np = np.array([15, -0.2, 1, 3, 0])
          self.r500_np = np.array([30, 1, 3, 30, 1])
          self.result = integrand_radial(r=self.r_np, r500=self.r500_np)
  
      def test_result(self):
          expectation = np.array([18.6780484, np.nan, np.nan, 19.9742601, np.nan])
          assert self.result == pytest.approx(expectation, nan_ok=True)
  
  
  class TestCorrectionRadial(TestCase):
  
      # TODO: add additional test cases.
  
      def setUp(self):
          # self.r_np = np.array([15, -0.2, 1, 3, 0])
          # self.r500_np = np.array([30, 1, 3, 30, 1])
          self.r500_np = 30
          self.ct = CorrectionRadial(
              integrand=integrand_radial, x_lim=self.r500_np)
  
      def test_result(self):
          """
  
          Integral determined "manually", split integral into two components,
          one where x < 1 and another with x > 1. x = R / R_s, change variables
          of eq 6 in Klein et al (2017) to x, change limits.
  
          dR = R_s * dx
  
          1 = C_rad \int R_s^2 (x / (x^2 -1)) * 2 * \pi * f(x)
  
          limits R_500, 0 - > 3, 0 because by default R_s = R_500 / 3
  
          https://www.wolframalpha.com/input/?i=integrate+(x%2F(x%5E2+-1)+)*+(
          +1+-+(2%2Fsqrt(1+-x%5E2))*arctanh(sqrt((1-x)%2F(
          x%2B1))))+with+respect+to+x+between+0+and+1
  
          https://www.wolframalpha.com/input/?i=integrate+(x%2F(x%5E2+-1)+)*+(
          +1+-+(2%2Fsqrt(x%5E2+-+1))*arctan(sqrt((x-1)%2F(
          x%2B1))))+with+respect+to+x+between+1+and+3
  
  
          Returns
          -------
  
          """
          assert self.ct.value == pytest.approx(1.893180E-3)
  
  
  class TestSchecterFunction(TestCase):
      def setUp(self):
          self.magnitude = np.array([20, 21.0, 21.0, 21.0, np.nan, 0.0])
          self.m_star = np.array([4, 3.0, 22.0, 3.0, 10, 1])
          self.phi_star = np.array([1, 1, 1, 1, 1, 1])
          self.alpha = np.array([-1, -1, -1, 0, -1, -1])
          self.lf = SchecterFunction(phi_star=self.phi_star, alpha=self.alpha)
  
      def test_m_star_min_m(self):
          expectation = np.array([-6.4, -7.2, 0.4, -7.2, np.nan, 0.4])
          result = self.lf.m_star_min_m(a=self.m_star, b=self.magnitude)
          assert result == pytest.approx(expectation, nan_ok=True)
  
      def test_partial(self):
          expectation = np.array([
              3.98107171e-07, 6.30957344e-08, 2.51188643, 6.30957344e-08, np.nan,
              2.51188643
          ])
          result = self.lf.partial(a=self.m_star, b=self.magnitude)
          assert result == pytest.approx(expectation, nan_ok=True)
  
      def test_integrand(self):
          expectation = np.array([
              9.21033671e-01, 9.21033979e-01, 7.47097468e-02, 5.81133154e-08,
              np.nan, 7.47097468e-02
          ])
          result = self.lf.integrand(m_star=self.m_star, magnitude=self.magnitude)
          assert result == pytest.approx(expectation, nan_ok=True)
  
      # def test_result(self):  #     expectation = np.array([18.6780484,
      #  np.nan, np.nan, 19.9742601, np.nan])  #     assert self.result ==  #
      #  pytest.approx(expectation, nan_ok=True)
  
  
  # class TestCorrectionCompleteness2(TestCase):
  #     #TODO: Test vectorised version
  #
  #     def setUp(self):
  #         self.c_lim = np.array([])
  #         self.magnitude = np.array([20, 21.0, 21.0, 21.0, np.nan, 0.0])
  #         self.m_star = np.array([4, 3.0, 22.0, 3.0, 10, 1])
  #         self.phi_star = np.array([1, 1, 1, 1, 1, 1])
  #         self.alpha = np.array([-1, -1, -1, 0, -1, -1])
  #
  #     def test_lower(self):
  #         expectation = np.array([-0.6, -1.6, 1.74, -1.6, 5.4, -3.6])
  #         assert self.ct.lower == pytest.approx(expectation)
  #
  #     def test_upper_numerator(self):
  #         assert self.ct.upper_numerator == pytest.approx([5.25, 4.25, 23.25,
  #                                                          4.25, 11.25, 2.25])
  #
  #     def test_upper_denominator(self):
  #         assert self.ct.upper_denominator == pytest.approx(5)
  #
  #     def test_integral_numerator(self):
  #         assert self.ct.integral_numerator == pytest.approx(0.866962)
  #
  #     def test_integral_denominator(self):
  #         assert self.ct.integral_denominator == pytest.approx(-1.28238E-32)
  #
  #     def test_value(self):
  #         #TODO: double check this
  #         assert self.ct.value == pytest.approx(-6.76057019E31)
  #
  #     def test_value_max(self):
  #         # When m*(z) + 1.25 < clim we set Ccmp =1.
  #         assert self.ct_max.value == pytest.approx(1)
  
  
  class TestCorrectionCompleteness(TestCase):
      def setUp(self):
          self.ct = CorrectionCompleteness(m_star=20, c_lim=5)
          self.ct_max = CorrectionCompleteness(m_star=3, c_lim=5)
  
      def test_lower(self):
          assert self.ct.lower == pytest.approx(15.4)
  
      def test_upper_numerator(self):
          assert self.ct.numerator_hi == pytest.approx(21.25)
  
      def test_upper_denominator(self):
          assert self.ct.denominator_hi == pytest.approx(5)
  
      def test_integral_numerator(self):
          """
          wolfram alpha text
          integrate  0.4ln10 * e^(-10^(0.4(20-x))) with respect to x between
          15.4 and 21.25
  
          Returns
          -------
  
          """
          assert self.ct.integral_numerator == pytest.approx(0.866962)
  
      def test_integral_denominator(self):
          """
          Wolfram alpha text:
          integrate  0.4ln10 * e^(-10^(0.4(20-x))) with respect to x between
          15.4 and 5
          Returns
          -------
  
          """
          assert self.ct.integral_denominator == pytest.approx(-1.28238E-32)
  
      def test_value(self):
          # TODO: double check this
          assert self.ct.value == pytest.approx(-6.76057019E31)
  
      def test_value_max(self):
          # When m*(z) + 1.25 < clim we set Ccmp =1.
          assert self.ct_max.value == pytest.approx(1)
  
  
  class TestRadialFilter(TestCase):
      def setUp(self):
          self.z = 0.5
          r500 = (Angle(0.0689, u.deg).to(u.arcsec) / 1.75).to(u.Mpc)
          self.app = (r500 / glo.cosmo.arcsec_per_kpc_proper(self.z)).to(u.Mpc)
          h = 0.7
          self.r_core = 1 * u.Mpc / h
  
      # # idx_bad = np.where((sigma_i < 0.05) & ((r_i / rs) < 1.5))[0]
      # idx_bad = ((r_i / rs) < 1.2) & (sigma_i < 0.05)
      # corr_filter_r =  / nfwweightNORM
      # return corr_filter_r
  
      def test_upper(self):
          rf = RadialFilter(
              r=np.arrange(0, 2, 0.1),
              rs=self.app / 3,
              r_core=self.r_core,
              r_lim=self.app)
          # rf.upper.
  
  
  class TestAreaTCL(TestCase):
      def setUp(self, test_guid='1_8055', dataset='old'):
          self.c = load_test_candidate(test_guid, dataset)
          # from regions import CircleSkyRegion
          # area = CircleSkyRegion(self.c.sky_opt, self.c.x_lim).area
          # mask_pkl = os.path.join(self.c.processed,
          #                         '{0}.mask.skyregions.pkl'.format(self.c))
          #
          # with open(mask_pkl, "rb") as mask_pkl_in:
          #     regions = pickle.load(mask_pkl_in)
          # mask_pkl_in.close()
  
      def test_A_tcl(self):
          print('foo')
          bg = self.c.A_bg()
          assert 1 == 1
  
  
  class TestLambdaMCFM(TestCase):
      def setUp(self, test_guid='1_62', dataset='new'):
          pass
          # self.c = load_test_candidate(test_guid, dataset)[0]
  
      def test_1_14908(self):
          pass
          # z, r = [], []
          # for i, f in enumerate(self.c.files_sources):
          #     if i > 10:
          #         z.append(f.z)
          #         r.append(richness(f, col=Str('griz')))
          # foo = load_test_candidate( '1_7427', 'new') +  load_test_candidate('1_3229', 'new' )
          # bar = [c for c in load_data('new') if all(f.exists for f in c.files_sources)]
          #
          # for i, c in enumerate(bar):
          #     if i < 100:
          #         # c = c[0]
          #         rich = [richness(f, col=Str('griz')) for f in c.files_sources]
          #         import matplotlib.pyplot as plt
          #         plt.clf()
          #         ax = plt.subplot(111)
          #         ax.scatter(glo.z_sweep, rich, label='{0.guid} -z: {0.z}, '
          #                                   'lamda_rm: {0.lambda_rm}'.format(c))
          #         # ax.scatter(x, 2 * np.pi * x * sigma/nfwweightNORM, label='r')
          #         # ax.scatter(r/rs, sigma, label='r')
          #         plt.legend()
          #         plt.savefig('test_{0}_rich_z.png'.format(c.guid))
          # # pass
  
      def test_spec(self):
          bar = [
              c for c in load_training_data(from_pkl=False, stack_galaxies=False)
              if c.file_sources_lambda.exists
          ]
          rich_mcmf = np.array(
              [richness(c.file_sources_lambda, col=Str('griz')) for c in bar])
  
          import matplotlib.pyplot as plt
          plt.clf()
          ax = plt.subplot(111)
          rich_rm = np.array([c.lambda_rm for c in bar])
          idx_good = np.where(rich_mcmf > 5)[0]
          ax.scatter(rich_rm[idx_good], rich_mcmf[idx_good])
          ax.plot(np.arange(20, 200, 1), np.arange(20, 200, 1))
          # ax.scatter(x, 2 * np.pi * x * sigma/nfwweightNORM, label='r')
          # ax.scatter(r/rs, sigma, label='r')
          ax.set_xscale("log", nonposx='clip')
          ax.set_yscale("log", nonposy='clip')
          ax.set_xlabel('LAMBDA_RM')
          ax.set_ylabel('LAMBDA_MCMF')
          plt.legend()
          plt.savefig('test_mcfm_rm.png')
  
  
  """
  
  
  """
  