import sys
#don't do this at home
sys.path.append("..")
from analytics import Swaption, Index, VolatilitySurface
from analytics.scenarios import run_swaption_scenarios
from pandas.tseries.offsets import BDay
import datetime
import numpy as np
import pandas as pd
from scipy.interpolate import SmoothBivariateSpline
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt

def plot_df(df, spread_shock, vol_shock):
    fig = plt.figure()
    ax = fig.gca(projection='3d')
    ## use smoothing spline on a finer grid
    f = SmoothBivariateSpline(df.vol_shock.values, df.spread_shock.values,
                              df.pv.values)
    xx, yy = np.meshgrid(vol_shock, spread_shock)
    surf = ax.plot_surface(xx, yy, f(vol_shock, spread_shock).T, cmap=cm.viridis)
    ax.set_xlabel("volatility shock")
    ax.set_ylabel("spread shock")
    ax.set_zlabel("PV")

trade_date = datetime.date(2017, 2, 23)
ig27 = Index.from_name("IG", 27, '5yr', trade_date=trade_date)
ig27.ref = 62
payer1 = Swaption(ig27, datetime.date(2017, 4, 19), 65)
payer2 = Swaption(ig27, datetime.date(2017, 5, 17), 72)
payer1.notional = 100e6
payer2.notional = 100e6
date_range = pd.bdate_range(trade_date, pd.Timestamp('2017-04-19') - BDay(), freq = '5B')
vol_shock = np.arange(-0.15, 0.3, 0.01)
spread_shock = np.arange(-0.2, 0.3, 0.01)
vs = VolatilitySurface("IG", 27, trade_date=trade_date)
vol_surface = vs[vs.list()[-1]]

df1 = run_swaption_scenarios(payer1, date_range, spread_shock, vol_shock, vol_surface,
                             params=['pv_black', 'delta'])
df2 = run_swaption_scenarios(payer2, date_range, spread_shock, vol_shock, vol_surface,
                             params=['pv_black', 'delta'])
#plot it
df = df1
df = df.assign(pv=df1.pv_black-df2.pv_black)
plot_df(df.loc[date_range[-1]], np.arange(-0.2, 0.3, 0.001),
        np.arange(-0.15, 0.3, 0.001))