from analytics.cms_spread import (
    quantlib_model, globeop_model, build_spread_index, VolatilityType,
    get_swaption_vol_data, get_swaption_vol_matrix)
from yieldcurve import YC
import pandas as pd
from quantlib.time.api import Date, Period, Years, Days, ModifiedFollowing
from quantlib.cashflows.cms_coupon import CmsCoupon
from quantlib.experimental.coupons.cms_spread_coupon import (
    CmsSpreadCoupon, CappedFlooredCmsSpreadCoupon)
import datetime
import math
# swap_index_30y2y, yc = build_spread_index(30, 2)
# cap = 0.00758
# corr = 0.8
# r = []
# maturity = pd.Timestamp("2020-01-19")
# today = pd.Timestamp.today()

# for d in pd.bdate_range("2018-01-19", today, closed="left", normalize=True):
#     d = pd.Timestamp("2018-01-19")
#     yc.link_to(YC(evaluation_date=d.date()))
#     yc.extrapolation = True
#     if d == pd.Timestamp("2018-02-16"):
#         continue
#     capped_floored_cms_spread_coupon_ln = \
#         quantlib_model(d, swap_index_30y2y, yc, cap, corr, maturity)
#     rate1 = capped_floored_cms_spread_coupon_ln.rate
#     cms_spread_coupon_n = quantlib_model(d, swap_index_30y2y, yc, cap, corr, maturity,
#                                          VolatilityType.Normal)
#     rate2 = cms_spread_coupon_n.rate
#     rate3 = globeop_model(d, swap_index_30y2y, yc, cap, corr - 0.075, maturity)
# # df = pd.DataFrame(r, columns=['date', 'QL_ln', 'QL_n', 'Globeop']).set_index('date')

# try to get convexity adjustement
from quantlib.indexes.swap.usd_libor_swap import UsdLiborSwapIsdaFixAm
from quantlib.cashflows.conundrum_pricer import AnalyticHaganPricer, YieldCurveModel
from quantlib.quotes import SimpleQuote
from quantlib.time.api import Actual365Fixed
# trade_date = datetime.date(2018, 1, 19)
trade_date = datetime.date(2018, 8, 23)
maturity = Date.from_datetime(trade_date) + Period(2, Years)
spread_index, yc = build_spread_index(30, 2)
yc.link_to(YC(evaluation_date=trade_date))
yc.extrapolation = True
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)

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

date, surf = get_swaption_vol_data(date=trade_date, vol_type=VolatilityType.ShiftedLognormal)
atm_vol = get_swaption_vol_matrix(trade_date, surf)
μ = SimpleQuote(0.1)
from quantlib.experimental.coupons.lognormal_cmsspread_pricer import LognormalCmsSpreadPricer
corr = SimpleQuote(0.8)
cms_pricer = AnalyticHaganPricer(atm_vol, YieldCurveModel.Standard, μ)
spread_pricer = LognormalCmsSpreadPricer(
    cms_pricer,
    corr,
    integration_points=20)
cms2y30y_cap.set_pricer(spread_pricer)
fixing_time = atm_vol.time_from_reference(cms2y30y_cap.fixing_date)
cms30y = CmsCoupon(payment_date,
                   100_000_000,
                   start_date=accrued_start_date,
                   end_date=accrued_end_date,
                   fixing_days=2,
                   index=spread_index.swap_index1,
                   is_in_arrears=True)
cms2y = CmsCoupon(payment_date,
                  100_000_000,
                  start_date=accrued_start_date,
                  end_date=accrued_end_date,
                  fixing_days=2,
                  index=spread_index.swap_index2,
                  is_in_arrears=True)
cms30y.set_pricer(cms_pricer)
cms2y.set_pricer(cms_pricer)
s1 = cms30y.index_fixing
s2 = cms2y.index_fixing
adjusted1 = cms30y.rate
adjusted2 = cms2y.rate
import math
T_alpha = atm_vol.time_from_reference(cms2y30y_cap.fixing_date)
mu1 = 1 / T_alpha * math.log(adjusted1 / s1)
mu2 = 1 / T_alpha * math.log(adjusted2 / s2)
vol1 = atm_vol.volatility(cms2y.fixing_date, spread_index.swap_index1.tenor, s1)
vol2 = atm_vol.volatility(cms30y.fixing_date, spread_index.swap_index2.tenor, s2)
mu_x = (mu1 - 0.5 * vol1 ** 2) * T_alpha
mu_y = (mu2 - 0.5 * vol2 ** 2) * T_alpha
sigma_x = vol1 * math.sqrt(T_alpha)
sigma_y = vol2 * math.sqrt(T_alpha)

from scipy.special import roots_hermitenorm
from analytics.cms_spread import h_call, h_put
import numpy as np
x, w = roots_hermitenorm(16)
val_put = 1/math.sqrt(2*math.pi) * np.dot(w, h_put(x, cap, s1, s2, mu_x, mu_y, sigma_x, sigma_y, corr.value))
val_call = 1/math.sqrt(2*math.pi) * np.dot(w, h_call(x, cap, s1, s2, mu_x, mu_y, sigma_x, sigma_y, corr.value))
print(cms2y30y_cap.rate, cms2y30y_cap.underlying.rate + val_put, val_call)