import unittest
from pyisda.cdsone import upfront_charge
from pyisda.curve import SpreadCurve
from pyisda.legs import ContingentLeg, FeeLeg
from pyisda.date import roll_date
import datetime
import numpy as np
from quantlib.settings import Settings
from quantlib.time.api import Date

import sys
sys.path.append('..')
from analytics import CreditIndex
from yieldcurve import YC, ql_to_jp, get_curve

class TestUpfront(unittest.TestCase):
    index = CreditIndex("ig", 26, "5yr",
                        value_date=datetime.date(2016, 9, 21))
    index.notional = 50e6
    index.spread = 70

    def test_upfront(self):
        self.assertAlmostEqual(-self.index.pv, 685292.81, 2)

    def test_cdsone(self):
        jp_yc = get_curve(self.index.value_date)
        fee_dirty = self.index.notional * (
            upfront_charge(self.index.value_date, self.index._cash_settle_date,
                           self.index.start_date,
                           self.index._step_in_date, self.index.start_date,
                           self.index.end_date, self.index.fixed_rate*1e-4,
                           jp_yc, 70e-4, self.index.recovery,
                           False))
        fee_clean = self.index.notional * (
            upfront_charge(self.index.value_date, self.index._cash_settle_date,
                           self.index.start_date,
                           self.index._step_in_date, self.index.start_date,
                           self.index.end_date, self.index.fixed_rate*1e-4,
                           jp_yc, 70e-4, self.index.recovery,
                           True))
        self.assertAlmostEqual(-fee_dirty, 685292.81, 2)
        self.assertAlmostEqual(fee_clean, self.index.clean_pv)

    def test_annuity(self):
        self.assertAlmostEqual(-self.index.clean_pv,
                               self.index.notional * self.index.risky_annuity *
                               (self.index.fixed_rate - self.index.spread)*1e-4)

class TestSpreadCurve(unittest.TestCase):

    def test_upfront_curves(self):
        upfront_curve = 1e-2 * np.array(
            [-2.502394, -4.871879, -9.329793, -12.98734, -15.833254,
             -17.622571, -20.505054, -24.314297])
        spread_curve = 1e-4 * np.array([500, 500, 500, 500, 500, 500, 500, 500])
        recovery_curve = np.full(8, 0.3)

        trade_date = datetime.date(2018, 6, 18)
        step_in_date = datetime.date(2018, 6, 19)
        cash_settle_date = datetime.date(2018, 6, 21)

        yc = get_curve(trade_date, "USD")
        tenors = np.array((0.5, 1, 2, 3, 4, 5, 7, 10))
        sc = SpreadCurve(trade_date, yc, None, None, None,
                         tenors, spread_curve, upfront_curve, recovery_curve, ticker="AES")
        maturities = [roll_date(trade_date, t) for t in tenors]
        for m, upf in zip(maturities, upfront_curve):
            pl = ContingentLeg(datetime.date(2017, 9, 20), m, 1)
            cl = FeeLeg(datetime.date(2017, 9, 20), m, True, 1, 0.05)
            a = pl.pv(trade_date, step_in_date, cash_settle_date, yc, sc, 0.3)
            b = cl.pv(trade_date, step_in_date, cash_settle_date, yc, sc, True)
            self.assertAlmostEqual(a - b, upf)

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