import sys
#don't do this at home
sys.path.append("..")
from analytics import Swaption, BlackSwaption, Index, VolatilitySurface
from analytics.scenarios import run_swaption_scenarios, run_index_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

import os
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import AxesGrid

def shiftedColorMap(cmap, start=0, midpoint=0.5, stop=1.0, name='shiftedcmap'):
    '''
    Function to offset the "center" of a colormap. Useful for
    data with a negative min and positive max and you want the
    middle of the colormap's dynamic range to be at zero

    Input
    -----
      cmap : The matplotlib colormap to be altered
      start : Offset from lowest point in the colormap's range.
          Defaults to 0.0 (no lower ofset). Should be between
          0.0 and `midpoint`.
      midpoint : The new center of the colormap. Defaults to
          0.5 (no shift). Should be between 0.0 and 1.0. In
          general, this should be  1 - vmax/(vmax + abs(vmin))
          For example if your data range from -15.0 to +5.0 and
          you want the center of the colormap at 0.0, `midpoint`
          should be set to  1 - 5/(5 + 15)) or 0.75
      stop : Offset from highets point in the colormap's range.
          Defaults to 1.0 (no upper ofset). Should be between
          `midpoint` and 1.0.
    '''
    cdict = {
        'red': [],
        'green': [],
        'blue': [],
        'alpha': []
    }

    # regular index to compute the colors
    reg_index = np.linspace(start, stop, 257)

    # shifted index to match the data
    shift_index = np.hstack([
        np.linspace(0.0, midpoint, 128, endpoint=False),
        np.linspace(midpoint, 1.0, 129, endpoint=True)
    ])

    for ri, si in zip(reg_index, shift_index):
        r, g, b, a = cmap(ri)

        cdict['red'].append((si, r, r))
        cdict['green'].append((si, g, g))
        cdict['blue'].append((si, b, b))
        cdict['alpha'].append((si, a, a))

    newcmap = matplotlib.colors.LinearSegmentedColormap(name, cdict)
    plt.register_cmap(cmap=newcmap)

    return newcmap

def plot_df(df, spread_shock, vol_shock, attr="pv"):
    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_color_map(df, spread_shock, vol_shock, attr="pv", path="."):

    val_date = df.index[0].date()
    #rows are spread, columns are volatility surface shift
    fig, ax = plt.subplots()
    series = df[attr]
    #Different ways to do a colormap: imshow and pcolormesh. using imshow here
    midpoint = 1 - series.max() / (series.max() + abs(series.min()))
    shifted_cmap = shiftedColorMap(cm.RdYlGn, midpoint=midpoint, name='shifted')

    chart = ax.imshow(series.values.reshape(spread_shock.size, vol_shock.size).T,
                      extent=(spread_shock.min(), spread_shock.max(),
                              vol_shock.min(), vol_shock.max()),
                      aspect='auto', interpolation='bilinear', cmap=shifted_cmap)

    ax.set_xlabel('Spread')
    ax.set_ylabel('Volatility shock')
    ax.set_title('{} of Trade on {}'.format(attr.title(), val_date))

    fig.colorbar(chart, shrink=.8)
    fig.savefig(os.path.join(path, "payer_swap_{}.png".format(val_date)))


trade_date = datetime.date(2017, 2, 23)
ig27 = Index.from_name("IG", 27, '5yr', trade_date=trade_date)
ig27.ref = 62
ig27.notional = 13e6
payer1 = BlackSwaption(ig27, datetime.date(2017, 4, 19), 65)
payer2 = BlackSwaption(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)
df2 = run_swaption_scenarios(payer2, date_range, spread_shock, vol_shock, vol_surface)

# #plot it
week = -1
df = df1
df = df.assign(pv=df1.pv-df2.pv)
spread_plot_range = ig27.ref * (1 + np.arange(-0.2, 0.3, 0.001))
vol_shock_range = np.arange(-0.15, 0.3, 0.001)
plot_df(df.loc[date_range[week]], spread_plot_range, vol_shock_range)
plot_color_map(df.loc[date_range[week]], ig27.ref * (1 + spread_shock), vol_shock)