import datetime
import math
import numpy as np
import unittest
from analytics.cms_spread import (
    build_spread_index, VolatilityType,
    get_swaption_vol_data, get_swaption_vol_matrix, get_cms_coupons,
    get_params, h_call, h_put)
from quantlib.quotes import SimpleQuote
from quantlib.time.api import (Actual365Fixed, Days, Date, ModifiedFollowing,
                               Period, Years)
from yieldcurve import YC
from quantlib.cashflows.conundrum_pricer import (
    AnalyticHaganPricer, YieldCurveModel)
from quantlib.experimental.coupons.cms_spread_coupon import (
    CappedFlooredCmsSpreadCoupon)
from quantlib.experimental.coupons.lognormal_cmsspread_pricer import (
    LognormalCmsSpreadPricer)
from scipy.special import roots_hermitenorm


class TestCmsSpread(unittest.TestCase):

    def setUp(self):
        self.trade_date = datetime.date(2018, 1, 19)
        option_tenor = Period(2, Years)
        maturity = Date.from_datetime(self.trade_date) + Period(2, Years)
        spread_index, self.yc = build_spread_index(30, 2)
        fixing_date = (spread_index.
                       fixing_calendar.
                       adjust(maturity, ModifiedFollowing))
        payment_date = (spread_index.
                        fixing_calendar.
                        advance(fixing_date, 2, Days))
        accrued_end_date = payment_date
        accrued_start_date = accrued_end_date - Period(1, Years)

        self.cap = 0.0075835

        self.notional = 100_000_000
        self.cms30y2y_cap = CappedFlooredCmsSpreadCoupon(
            payment_date,
            self.notional,
            start_date=accrued_start_date,
            end_date=accrued_end_date,
            fixing_days=spread_index.fixing_days,
            index=spread_index,
            gearing=1.,
            spread=-self.cap,
            floor=0.,
            day_counter=Actual365Fixed(),
            is_in_arrears=True)

        self.cms2y, self.cms30y = get_cms_coupons(self.trade_date,
                                                  self.notional,
                                                  option_tenor,
                                                  spread_index)

    def test_black_model(self):
        evaluation_date = datetime.date(2018, 8, 23)
        self.yc.link_to(YC(evaluation_date=evaluation_date))
        self.yc.extrapolation = True
        date, surf = get_swaption_vol_data(date=evaluation_date,
                                           vol_type=VolatilityType.ShiftedLognormal)
        atm_vol = get_swaption_vol_matrix(evaluation_date, surf)
        μ = SimpleQuote(0.1)
        ρ = SimpleQuote(0.8)
        cms_pricer = AnalyticHaganPricer(atm_vol, YieldCurveModel.Standard, μ)
        self.cms2y.set_pricer(cms_pricer)
        self.cms30y.set_pricer(cms_pricer)
        params = get_params(self.cms2y, self.cms30y, atm_vol)
        x, w = roots_hermitenorm(16)
        val_call = 1 / math.sqrt(2 * math.pi) * np.dot(w,
                                                       h_call(x, self.cap, *params, ρ.value))
        val_put = 1 / math.sqrt(2 * math.pi) * np.dot(w,
                                                      h_put(x, self.cap, *params, ρ.value))
        cms_spread_pricer = LognormalCmsSpreadPricer(
            cms_pricer,
            ρ,
            integration_points=20)
        self.cms30y2y_cap.set_pricer(cms_spread_pricer)
        self.assertAlmostEqual(self.cms30y2y_cap.rate,
                               val_call)
        self.assertAlmostEqual(-self.cms30y2y_cap.underlying.rate, val_put - val_call)
        self.assertAlmostEqual(self.cms30y.rate - self.cms2y.rate - self.cap,
                               self.cms30y2y_cap.underlying.rate)

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