import unittest
import numpy as np
from pyisda.legs import ContingentLeg, FeeLeg
from pyisda.curve import SpreadCurve
from pickle import dumps, loads
import datetime
from pyisda.utils import build_yc


class TestPickle(unittest.TestCase):
    trade_date = datetime.date(2016, 10, 6)
    yc = build_yc(trade_date, ql_curve=True)

    def assertListAlmostEqual(self, l1, l2):
        if len(l1) != len(l2):
            raise AssertionError
        for a, b in zip(l1, l2):
            self.assertAlmostEqual(a, b)

    def test_yc(self):
        orig_dfs = [self.yc.discount_factor(d) for d in self.yc.dates]
        pickled_yc = loads(dumps(self.yc, 2))
        self.assertListAlmostEqual(
            [pickled_yc.discount_factor(d) for d in pickled_yc.dates], orig_dfs
        )

    def test_legs(self):
        start_date = datetime.date(2016, 3, 20)
        end_date = datetime.date(2021, 12, 20)
        step_in_date = datetime.date(2016, 10, 7)
        value_date = datetime.date(2016, 10, 11)
        cl = ContingentLeg(start_date, end_date, 1)
        fl = FeeLeg(start_date, end_date, True, 1, 1)
        sc = SpreadCurve(
            self.trade_date,
            self.yc,
            start_date,
            step_in_date,
            value_date,
            [end_date],
            np.array([75 * 1e-4]),
            0.4,
        )
        sc_pickled = loads(dumps(sc, 2))
        self.assertListAlmostEqual(
            [sc.survival_probability(d) for d in self.yc.dates],
            [sc_pickled.survival_probability(d) for d in self.yc.dates],
        )
        cl_pickled = loads(dumps(cl, 2))
        fl_pickled = loads(dumps(fl, 2))
        self.assertEqual(
            cl_pickled.pv(self.trade_date, step_in_date, value_date, self.yc, sc, 0.4),
            cl.pv(self.trade_date, step_in_date, value_date, self.yc, sc, 0.4),
        )
        self.assertEqual(
            fl.pv(self.trade_date, step_in_date, value_date, self.yc, sc, False),
            fl_pickled.pv(
                self.trade_date, step_in_date, value_date, self.yc, sc, False
            ),
        )


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