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 serenitas.analytics import CreditIndex
from serenitas.analytics.yieldcurve import 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)

    @staticmethod
    def aux(index, yc, clean):
        """computes upfront fee from index using upfront_charge of pyisda"""
        return index.notional * (
            upfront_charge(
                index.value_date,
                index._cash_settle_date,
                index.start_date,
                index._step_in_date,
                index.start_date,
                index.end_date,
                index.fixed_rate * 1e-4,
                yc,
                index.spread * 1e-4,
                index.recovery,
                clean,
            )
        )

    def test_cdsone(self):
        jp_yc = get_curve(self.index.value_date, "USD")
        fee_dirty = self.aux(self.index, jp_yc, False)
        fee_clean = self.aux(self.index, jp_yc, True)
        self.assertAlmostEqual(-fee_dirty, 685292.81, 2)
        self.assertAlmostEqual(fee_clean, self.index.clean_pv)

    def test_ois(self):
        index = CreditIndex("HY", 37, "5yr", value_date=datetime.date(2022, 2, 4))
        index.notional = 100_000_000
        index.spread = 1000
        yc_ois = get_curve(index.value_date, "USD", True)
        yc = get_curve(index.value_date, "USD")
        fee_clean_ois = self.aux(index, yc_ois, True)
        fee_clean = self.aux(index, yc, True)
        self.assertAlmostEqual(17154197.82, fee_clean, 2)  # bloomberg curve 260
        self.assertAlmostEqual(17229522.44, fee_clean_ois, 2)  # bloomberg curve 531

    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 setUp(self):
        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)

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

        self.yc = get_curve(self.trade_date, "USD")
        self.tenors = np.array((0.5, 1, 2, 3, 4, 5, 7, 10))
        self.sc = SpreadCurve(
            self.trade_date,
            self.yc,
            None,
            None,
            None,
            self.tenors,
            spread_curve,
            self.upfront_curve,
            recovery_curve,
            ticker="AES",
        )

    def test_upfront_curves(self):

        maturities = [roll_date(self.trade_date, t) for t in self.tenors]
        for m, upf in zip(maturities, self.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(
                self.trade_date,
                self.step_in_date,
                self.cash_settle_date,
                self.yc,
                self.sc,
                0.3,
            )
            b = cl.pv(
                self.trade_date,
                self.step_in_date,
                self.cash_settle_date,
                self.yc,
                self.sc,
                True,
            )
            self.assertAlmostEqual(a - b, upf)

    def test_roundtrip(self):
        sc_copy = SpreadCurve.from_bytes(self.sc.as_buffer(True), True)
        self.assertEqual(sc_copy.inspect(), self.sc.inspect())
        self.assertEqual(sc_copy.full_ticker, self.sc.full_ticker)
        sc_copy = SpreadCurve.from_bytes(self.sc.as_buffer(False), False)
        self.assertEqual(sc_copy.inspect(), self.sc.inspect())
        self.assertEqual(sc_copy.full_ticker, self.sc.full_ticker)


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