import unittest
import array
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))
        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], array.array('d', [75*1e-4]), 0.4)
        sc_pickled = loads(dumps(sc))
        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))
        fl_pickled = loads(dumps(fl))
        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))