import sys
#don't do this at home
sys.path.append("..")
from analytics import (Swaption, BlackSwaption, BlackSwaptionVolSurface,
                       Index, ProbSurface, Portfolio)
from analytics.scenarios import run_swaption_scenarios, run_index_scenarios, run_portfolio_scenarios
from pandas.tseries.offsets import BDay
import datetime
import numpy as np
import pandas as pd
from scipy.interpolate import SmoothBivariateSpline

import os
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from graphics import plot_time_color_map, plot_color_map

from db import dbengine
engine  = dbengine('serenitasdb')

def plot_df(df, spread_shock, vol_shock, attr="pnl"):
    val_date = df.index[0].date()
    fig = plt.figure()

    ax = fig.gca(projection='3d')
    ## use smoothing spline on a finer grid
    series = df[attr]
    f = SmoothBivariateSpline(df.vol_shock.values, df.spread_shock.values, series.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")
    ax.set_zlabel("PnL")
    ax.set_title('{} of Trade on {}'.format(attr.title(), val_date))

def plot_trade_scenarios(portf, shock_min=-.15, shock_max=.2, period=-1, vol_time_roll=True):
    portf.reset_pv()
    earliest_date = min(portf.swaptions, key=lambda x: x.exercise_date).exercise_date
    date_range = pd.bdate_range(portf.indices[0].trade_date,
                                earliest_date - BDay(), freq='3B')
    vol_shock = np.arange(-0.15, 0.3, 0.01)
    spread_shock = np.arange(shock_min, shock_max, 0.01)
    index = portf.indices[0].name.split()[1]
    series = portf.indices[0].name.split()[3][1:]
    vs = BlackSwaptionVolSurface(index, series, trade_date=portf.indices[0].trade_date)
    vol_surface = vs[vs.list(option_type='payer')[-1]]

    df = run_portfolio_scenarios(portf, date_range, spread_shock, vol_shock, vol_surface,
                                 params=["pnl","delta"])

    hy_plot_range = 100 + (500 - portf.indices[0].spread * (1 + spread_shock)) * \
                    abs(portf.indices[0].DV01) / portf.indices[0].notional * 100

    shock =  hy_plot_range if index == 'HY' else portf.indices[0].spread * (1 + spread_shock)

    plot_time_color_map(df[round(df.vol_shock,2)==0], shock, 'pnl', index=index)
    plot_time_color_map(df[round(df.vol_shock,2)==.2], shock, 'pnl', index=index)
    plot_color_map(df.loc[date_range[period]], shock, vol_shock, 'pnl', index=index)
    return df

def exercise_probability():
    engine  = dbengine('serenitasdb')
    #Ad hoc
    option_delta = Index.from_name('HY', 29, '5yr')
    option_delta.price = 107.875
    option1 = BlackSwaption(option_delta, datetime.date(2017, 12, 20), 107, option_type="payer")
    option2 = BlackSwaption(option_delta, datetime.date(2017, 12, 20), 105, option_type="payer")
    option1.sigma = .280
    option2.sigma = .371
    option1.notional = 20_000_000
    option2.notional = 40_000_000
    option1.direction = 'Long'
    option2.direction = 'Short'
    option_delta.notional = option1.notional * option1.delta + option2.notional * option2.delta
    option_delta.direction = 'Seller' if option_delta.notional > 0 else 'Buyer'
    option_delta.notional = abs(option_delta.notional)
    portf = Portfolio([option1, option2, option_delta])

    portf.reset_pv()
    earliest_date = min(portf.swaptions, key=lambda x: x.exercise_date).exercise_date
    date_range = pd.bdate_range(portf.indices[0].trade_date, earliest_date - BDay(), freq='5B')
    vol_shock = np.arange(-0.15, 0.3, 0.01)
    spread_shock = np.arange(-0.15, 0.35, 0.01)
    index = portf.indices[0].name.split()[1]
    series = portf.indices[0].name.split()[3][1:]

    vs = ProbSurface(index, series, trade_date=portf.indices[0].trade_date)
    vs.plot(vs.list()[-1])