from __future__ import division, print_function
import numpy

from foldedleastsquares import DefaultTransitTemplateGenerator

import unittest

class TestPeriodGrid(unittest.TestCase):
    def test(self):
        print("Starting test: period_grid...", end="")
        default_transit_template_generator = DefaultTransitTemplateGenerator()
        periods = default_transit_template_generator.period_grid(R_star=1, M_star=1, time_span=0.1)
        numpy.testing.assert_almost_equal(max(periods), 2.4999999999999987)
        numpy.testing.assert_almost_equal(min(periods), 0.6002621413799498)
        numpy.testing.assert_equal(len(periods), 268)

        periods = default_transit_template_generator.period_grid(R_star=1, M_star=1, time_span=20)
        numpy.testing.assert_almost_equal(max(periods), 10)
        numpy.testing.assert_almost_equal(min(periods), 0.6015575922909607)
        numpy.testing.assert_equal(len(periods), 1716)

        periods = default_transit_template_generator.period_grid(
            R_star=5,  # R_sun
            M_star=1,  # M_sun
            time_span=20,  # days
            period_min=0,
            period_max=999,
            oversampling_factor=3,
        )
        numpy.testing.assert_almost_equal(max(periods), 10)
        numpy.testing.assert_almost_equal(min(periods), 0.6015575922909607)
        numpy.testing.assert_equal(len(periods), 1716)

        periods = default_transit_template_generator.period_grid(
            R_star=1,  # R_sun
            M_star=1,  # M_sun
            time_span=20,  # days
            period_min=0,
            period_max=999,
            oversampling_factor=3,
        )
        numpy.testing.assert_almost_equal(max(periods), 10)
        numpy.testing.assert_almost_equal(min(periods), 0.60155759)
        numpy.testing.assert_equal(len(periods), 1716)

        periods = default_transit_template_generator.period_grid(
            R_star=0.1,  # R_sun
            M_star=1,  # M_sun
            time_span=1000,  # days
            period_min=0,
            period_max=999,
            oversampling_factor=3,
        )
        numpy.testing.assert_equal(len(periods), 4308558)

        periods = default_transit_template_generator.period_grid(
            R_star=1,  # R_sun
            M_star=1,  # M_sun
            time_span=20,  # days
            period_min=0.5,
            period_max=999,
            oversampling_factor=3,
        )
        numpy.testing.assert_almost_equal(max(periods), 10)
        numpy.testing.assert_almost_equal(min(periods), 0.500, 3)
        numpy.testing.assert_equal(len(periods), 1895)
        print("passed")
