import unittest
import datetime

from serenitas.analytics.index import CreditIndex, ForwardIndex
from serenitas.analytics.basket_index import BasketIndex
from serenitas.analytics.index import g
from serenitas.utils.db import dbconn
import pickle


class TestPickle(unittest.TestCase):
    index = CreditIndex("ig", 26, "5yr", value_date=datetime.date(2016, 7, 1))
    index.notional = 50e6
    index.spread = 75

    def test_pickle(self):
        a = pickle.loads(pickle.dumps(self.index))
        self.assertTrue(hash(a) == hash(self.index))

    def test_pickle_basket(self):
        a = BasketIndex("IG", 31, ["5yr"], value_date=datetime.date(2019, 2, 19))
        self.assertTrue(hash(pickle.loads(pickle.dumps(a))), hash(a))

    def test_from_tradeid(self):
        ig28 = CreditIndex.from_tradeid(874)
        self.assertTrue(ig28.spread, 68.0)


class TestStrike(unittest.TestCase):
    index = CreditIndex("ig", 26, "5yr", value_date=datetime.date(2016, 7, 1))
    index.notional = 50_000_000.0
    index.spread = 75
    exercise_date = datetime.date(2016, 8, 19)

    def test_pv(self):
        self.assertAlmostEqual(
            self.index.clean_pv,
            g(self.index, self.index.spread, self.index.value_date)
            * self.index.notional,
        )

    def test_strike(self):
        """strike price equals clean_pv using expected forward yield curve"""
        strike = (
            g(self.index, self.index.spread, self.exercise_date) * self.index.notional
        )
        old_yc = self.index._yc
        self.index.value_date = self.exercise_date
        self.index._yc = old_yc.expected_forward_curve(self.exercise_date)
        self.index._update_spread_curve()
        self.index._update_pvs()
        self.assertAlmostEqual(self.index.clean_pv, strike)

    def test_price_setting(self):
        self.index.price = self.index.price
        self.assertAlmostEqual(self.index.spread, 75)

    def test_pv_setting(self):
        self.index.pv = self.index.pv
        self.assertAlmostEqual(self.index.spread, 75)
        self.index.direction = "Seller"
        self.index.pv = self.index.pv
        self.assertAlmostEqual(self.index.spread, 75)


class TestForwardIndex(unittest.TestCase):
    index = CreditIndex("ig", 26, "5yr", value_date=datetime.date(2016, 7, 1))
    index.notional = 50_000_000.0
    index.spread = 75
    exercise_date = datetime.date(2016, 8, 19)
    fi = ForwardIndex(index, exercise_date)

    def test_forward_annuity(self):
        """default adjusted forward spread and forward annuity match"""
        # TODO: check what's going on
        self.assertAlmostEqual(
            self.fi.forward_pv,
            self.fi.forward_annuity
            * (self.fi.forward_spread - self.fi.index.fixed_rate)
            * 1e-4,
            4,
        )

    def test_forward_pv(self):
        """default adjusted forward price for trade_date equals clean pv"""
        fi = ForwardIndex(self.index, self.index.value_date)
        self.assertAlmostEqual(fi.forward_pv, self.index.clean_pv / self.index.notional)


class TestIndexTrade(unittest.TestCase):
    def test_traded_level(self):
        trade = CreditIndex.from_tradeid(3436)
        with trade._trade_conn.cursor() as c:
            c.execute("SELECT traded_level FROM cds WHERE id=%s", (3436,))
            (spread,) = c.fetchone()
        self.assertAlmostEqual(trade.spread, float(spread))

    def test_sofr_trade(self):
        trade = CreditIndex.from_tradeid(3672)
        with trade._trade_conn.cursor() as c:
            c.execute("SELECT traded_level FROM cds WHERE id=%s", (3672,))
            (spread,) = c.fetchone()
        self.assertAlmostEqual(trade.spread, float(spread), 6)


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