import datetime
from analytics.curve_trades import curve_pos, on_the_run
from analytics.index_data import index_returns
import numpy as np
import pandas as pd

index_type = "IG"
portf = curve_pos(datetime.date(2018, 5, 3), index_type)

df = index_returns(index=index_type, years=5,
                   tenor=['3yr', '5yr', '7yr', '10yr'])
df = (df.reset_index(['index'], drop=True).
      reorder_levels(['date', 'series', 'tenor']))
returns = df.spread_return.dropna().reset_index('series')
returns['dist_on_the_run'] = (returns.
                              groupby('date')['series'].
                              transform(lambda x: x.max() - x))
del returns['series']
returns = returns.set_index('dist_on_the_run', append=True).unstack('tenor')
returns.columns = returns.columns.droplevel(0)
portf.reset_pv()

otr = on_the_run(index_type)
spreads = pd.DataFrame({'spread': portf.spread,
                        'tenor': [ind.tenor for ind in portf.indices],
                        'dist_on_the_run': [otr - ind.series for ind in portf.indices]})
spreads = spreads.set_index(['dist_on_the_run', 'tenor'])
r = []
for k, g in returns.groupby(level='date', as_index=False):
    shocks = g.reset_index('date', drop=True).stack('tenor')
    shocks.name = 'shocks'
    portf.spread = spreads.spread * (1 + spreads.join(shocks).shocks)
    r.append((k, portf.pnl))
pnl = pd.DataFrame.from_records(r, columns=['date', 'pnl'], index=['date'])