import unittest
import array
from pyisda.legs import ContingentLeg, FeeLeg
from pyisda.curve import SpreadCurve
from pickle import dumps, loads
import datetime
from quantlib.settings import Settings
from quantlib.time.api import Date
import sys
sys.path.append('/home/share/CorpCDOs/code/python/')
from yieldcurve import YC, get_dates, ql_to_jp

class TestPickle(unittest.TestCase):

    settings = Settings()
    trade_date = datetime.date(2016, 10, 6)
    settings.evaluation_date =  Date(6, 10, 2016)
    yc = YC()
    jp_yc = ql_to_jp(yc)

    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):
        jp_yc = ql_to_jp(self.yc)
        dates = get_dates(self.trade_date)
        self.assertListAlmostEqual([jp_yc.discount_factor(d) for d in dates],
                                   self.yc.data[1:])
        pickled_yc = loads(dumps(jp_yc))
        self.assertListAlmostEqual([pickled_yc.discount_factor(d) for d in dates],
                                   self.yc.data[1:])

    def test_contingentleg(self):
        start_date = datetime.date(2016, 9, 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.jp_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 get_dates(self.trade_date)],
                                   [sc_pickled.survival_probability(d) for d in get_dates(self.trade_date)])
        cl_pickled = loads(dumps(cl))
        fl_pickled = loads(dumps(fl))
        self.assertEqual(cl_pickled.pv(self.trade_date, step_in_date, value_date, self.jp_yc, sc, 0.4),
                         cl.pv(self.trade_date, step_in_date, value_date, self.jp_yc, sc, 0.4))
        self.assertEqual(fl.pv(self.trade_date, step_in_date, value_date, self.jp_yc, sc, False),
                         fl_pickled.pv(self.trade_date, step_in_date, value_date, self.jp_yc, sc, False))