import unittest
from quantlib.time.api import (
    Date,
    Period,
    Quarterly,
    Schedule,
    Rule,
    WeekendsOnly,
    Following,
    Unadjusted,
)
from quantlib.settings import Settings
from serenitas.analytics.yieldcurve import YC


class TestYieldCurve(unittest.TestCase):
    def assertListAlmostEqual(self, list1, list2, places=7):
        self.assertEqual(len(list1), len(list2))
        for a, b in zip(list1, list2):
            self.assertAlmostEqual(a, b, places)

    def test_bloomberg(self):
        discounts = [
            0.99848606,
            0.99548212,
            0.99167201,
            0.98772518,
            0.98629694,
            0.98509013,
            0.98389804,
            0.98268094,
            0.97993369,
            0.97709955,
            0.97430462,
            0.97145583,
            0.96827211,
            0.9651054,
            0.96194904,
            0.95869946,
            0.95499459,
            0.95125679,
            0.94753361,
            0.94382501,
        ]
        settings = Settings()
        settings.evaluation_date = Date(23, 6, 2016)
        curve = YC()
        term_date = Date(20, 6, 2021)
        cds_schedule = Schedule.from_rule(
            settings.evaluation_date,
            term_date,
            Period(Quarterly),
            WeekendsOnly(),
            Following,
            Unadjusted,
            Rule.CDS2015,
        )
        curve_df = [curve.discount(d) for d in cds_schedule[1:-1]]
        # last date is term_date+1
        curve_df.append(curve.discount(term_date + 1))
        self.assertListAlmostEqual(discounts, curve_df)


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