import unittest
import datetime
import pandas as pd

from serenitas.analytics.utils import roll_date, days_accrued
from serenitas.analytics.tranche_functions import cds_accrued, credit_schedule
from pyisda.date import (
    roll_date as roll_date_c,
    cds_accrued as cds_accrued_c,
    previous_twentieth,
)
from serenitas.analytics.yieldcurve import get_curve


class TestStartDate(unittest.TestCase):
    def test_previous_twentieth(self):
        self.assertEqual(
            previous_twentieth(datetime.date(2016, 9, 20)), datetime.date(2016, 9, 20)
        )
        self.assertEqual(
            previous_twentieth(datetime.date(2016, 6, 10)), datetime.date(2016, 3, 21)
        )
        self.assertEqual(
            previous_twentieth(datetime.date(2016, 6, 10), False),
            datetime.date(2016, 3, 20),
        )
        self.assertEqual(
            previous_twentieth(datetime.date(2016, 10, 15)), datetime.date(2016, 9, 20)
        )
        self.assertEqual(
            previous_twentieth(datetime.date(2016, 12, 20)), datetime.date(2016, 12, 20)
        )
        self.assertEqual(
            previous_twentieth(datetime.date(2017, 3, 19)), datetime.date(2016, 12, 20)
        )

    def test_previous_twentieth_timestamp(self):
        self.assertEqual(
            previous_twentieth(pd.Timestamp("2016-09-20")), datetime.date(2016, 9, 20)
        )
        self.assertEqual(
            previous_twentieth(pd.Timestamp("2016-10-15")), datetime.date(2016, 9, 20)
        )
        self.assertEqual(
            previous_twentieth(pd.Timestamp("2016-12-20")), datetime.date(2016, 12, 20)
        )
        self.assertEqual(
            previous_twentieth(pd.Timestamp("2017-03-19")), datetime.date(2016, 12, 20)
        )


class TestEndDate(unittest.TestCase):
    def test_enddate(self):
        self.assertEqual(
            roll_date(datetime.date(2016, 9, 20), 5), datetime.date(2021, 12, 20)
        )
        self.assertEqual(
            roll_date(datetime.date(2016, 10, 15), 5), datetime.date(2021, 12, 20)
        )
        self.assertEqual(
            roll_date(datetime.date(2017, 3, 19), 5), datetime.date(2021, 12, 20)
        )
        self.assertEqual(
            roll_date(datetime.date(2017, 3, 20), 5), datetime.date(2022, 6, 20)
        )

    def test_enddate_pre2015(self):
        self.assertEqual(
            roll_date(datetime.date(2014, 9, 20), 5), datetime.date(2019, 12, 20)
        )
        self.assertEqual(
            roll_date(datetime.date(2014, 10, 15), 5), datetime.date(2019, 12, 20)
        )
        self.assertEqual(
            roll_date(datetime.date(2015, 3, 19), 5), datetime.date(2020, 3, 20)
        )
        self.assertEqual(
            roll_date(datetime.date(2015, 3, 20), 5), datetime.date(2020, 6, 20)
        )
        self.assertEqual(
            roll_date(datetime.date(2015, 12, 19), 5), datetime.date(2020, 12, 20)
        )
        self.assertEqual(
            roll_date(datetime.date(2015, 12, 20), 5), datetime.date(2020, 12, 20)
        )


class TestEndDateC(unittest.TestCase):
    def test_enddate(self):
        self.assertEqual(
            roll_date_c(datetime.date(2016, 9, 20), 5), datetime.date(2021, 12, 20)
        )
        self.assertEqual(
            roll_date_c(datetime.date(2016, 10, 15), 5), datetime.date(2021, 12, 20)
        )
        self.assertEqual(
            roll_date_c(datetime.date(2017, 3, 19), 5), datetime.date(2021, 12, 20)
        )
        self.assertEqual(
            roll_date_c(datetime.date(2017, 3, 20), 5), datetime.date(2022, 6, 20)
        )

    def test_enddate_pre2015(self):
        self.assertEqual(
            roll_date_c(datetime.date(2014, 9, 20), 5), datetime.date(2019, 12, 20)
        )
        self.assertEqual(
            roll_date_c(datetime.date(2014, 10, 15), 5), datetime.date(2019, 12, 20)
        )
        self.assertEqual(
            roll_date_c(datetime.date(2015, 3, 19), 5), datetime.date(2020, 3, 20)
        )
        self.assertEqual(
            roll_date_c(datetime.date(2015, 3, 20), 5), datetime.date(2020, 6, 20)
        )
        self.assertEqual(
            roll_date_c(datetime.date(2015, 12, 19), 5), datetime.date(2020, 12, 20)
        )
        self.assertEqual(
            roll_date_c(datetime.date(2015, 12, 20), 5), datetime.date(2020, 12, 20)
        )


class TestCdsAccrued(unittest.TestCase):
    def test_cdsaccrued(self):
        self.assertAlmostEqual(
            cds_accrued(datetime.date(2016, 4, 1), 100),
            cds_accrued_c(datetime.date(2016, 4, 1), 100),
        )
        self.assertAlmostEqual(
            cds_accrued(datetime.date(2016, 3, 21), 100), 0.2777777777777778
        )
        self.assertAlmostEqual(
            cds_accrued(datetime.date(2016, 3, 18), 100), 24.72222222222222
        )
        # buggy
        # self.assertAlmostEqual(cds_accrued(datetime.date(2016, 9, 19), 100),
        #                        0.)

    def test_cdsaccrued_c(self):
        self.assertAlmostEqual(
            cds_accrued_c(datetime.date(2016, 4, 1), 100), 3.3333333333333
        )
        self.assertAlmostEqual(
            cds_accrued_c(datetime.date(2016, 3, 21), 100), 0.27777777777777
        )
        self.assertAlmostEqual(
            cds_accrued_c(datetime.date(2016, 3, 18), 100), 24.72222222222
        )
        self.assertAlmostEqual(cds_accrued_c(datetime.date(2016, 9, 19), 100), 0.0)
        self.assertAlmostEqual(
            cds_accrued_c(datetime.date(2014, 9, 19), 100), 25.55555555556
        )


class TestDaysAccrued(unittest.TestCase):
    def test_days_accrued(self):
        dates = [
            ("2015-01-07", 17),
            ("2015-06-19", 92),
            ("2015-03-19", 0),
            ("2015-06-20", 93),
            ("2015-06-21", 0),
            ("2015-12-21", 1),
            ("2015-12-28", 8),
        ]
        for date, days in dates:
            self.assertEqual(days_accrued(date), days)


class TestCdsSchedule(unittest.TestCase):
    def setUp(self):
        self.trade_date = datetime.date(2020, 6, 19)
        self.yc = get_curve(self.trade_date, "USD")

    def test_credit_schedule(self):
        """ tests when IMM date falls on a week-end"""
        df = credit_schedule(self.trade_date, 1.0, self.yc, datetime.date(2022, 12, 20))
        self.assertEqual(df.index[0], pd.Timestamp("2020-06-22"))
        df = credit_schedule(
            datetime.date(2020, 3, 19), 1.0, self.yc, datetime.date(2022, 12, 20)
        )
        self.assertEqual(df.index[0], pd.Timestamp("2020-06-22"))


if __name__ == "__main__":
    unittest.main()