from analytics.index_data import get_index_quotes, index_returns
from db import dbengine
from analytics import Index, Portfolio
from analytics.utils import roll_date
from dateutil.relativedelta import relativedelta
import pandas as pd
import math
import statsmodels.formula.api as smf
from analytics.basket_index import MarkitBasketIndex
import numpy as np
import matplotlib.pyplot as plt

from statsmodels.sandbox.regression.predstd import wls_prediction_std

_engine = dbengine('serenitasdb')

def on_the_run(index):
    r = _engine.execute("SELECT max(series) FROM index_version WHERE index=%s",
                                  (index,))
    series, = r.fetchone()
    return series

def curve_spread_diff(index='IG', rolling=6):
    otr = on_the_run(index)
    ## look at spreads
    df = get_index_quotes(index, list(range(otr - rolling, otr + 1)),
                          tenor=['3yr', '5yr', '7yr', '10yr'])
    spreads = df.groupby(level=['date', 'tenor']).nth(-1)['closespread'].unstack(-1)
    spreads_diff = spreads.diff(axis=1)
    del spreads_diff['3yr']
    spreads_diff.columns = ['3-5', '5-7', '7-10']
    spreads_diff['5-10'] = spreads_diff['5-7'] + spreads_diff['7-10']
    return spreads_diff

def spreads_diff_table(spreads_diff):
    def current(s):
        return s.iat[-1]
    def zscore(s):
        return (s.iat[-1] - s.mean()) / s.std()
    df = spreads_diff.agg(['min', 'max','mean', current, zscore])
    ((spreads_diff - spreads_diff.mean())/spreads_diff.std()).plot()
    return df

def theta_matrix_by_series(index='IG', rolling=6):
    otr = on_the_run(index)
    df = get_index_quotes(index, list(range(otr - rolling, otr + 1)),
                          tenor=['3yr', '5yr', '7yr', '10yr'])
    df['theta_per_dur'] = df.theta2 / df.duration2
    theta_matrix = df.groupby(level=['date', 'tenor','series']).nth(-1)['theta_per_dur']
    theta_matrix = theta_matrix.loc[theta_matrix.index[-1][0]].unstack(0)
    return theta_matrix[['3yr', '5yr', '7yr', '10yr']]

def ratio_within_series(index='IG', rolling=6, param='duration'):
    otr = on_the_run(index)
    df = get_index_quotes(index, list(range(otr - rolling, otr + 1)),
                          tenor=['3yr', '5yr', '7yr', '10yr']).unstack()
    ratio = (df[param].
             apply(lambda s: s / df[param]['5yr'].values, raw=True))
    ratio.columns = pd.MultiIndex.from_product([[param + '_ratio_to_5yr'], ratio.columns])
    df = df.join(ratio).groupby(['date']).tail(1)
    df = df.reset_index(level=['index', 'version'], drop=True)
    return df

def curve_3_5_10(df):
    """
    Parameters
    ----------
    df: duration ratio within series"""
    #buy 3y, sell 5y, buy 10y
    s = - df.theta2['3yr'] / df.duration_ratio_to_5yr['3yr'] \
        + 2 * df.theta2['5yr'] \
        - df.theta2['10yr'] / df.duration_ratio_to_5yr['10yr']
    s.dropna().unstack(-1).plot()

def curve_5_10(df):
    #buy sell 5y, buy 10y
    s = df.theta2['5yr'] - df.theta2['10yr'] / df.duration_ratio_to_5yr['10yr']
    s.dropna().unstack(-1).plot()

def on_the_run_theta(index='IG', rolling=6):
    otr = on_the_run(index)
    df = get_index_quotes(index, list(range(otr - rolling, otr + 1)),
                          tenor=['3yr', '5yr', '7yr', '10yr'])
    df['theta_per_dur'] = df.theta2/df.duration2
    theta_matrix = df.groupby(level=['date', 'tenor']).nth(-1)['theta_per_dur']
    theta_matrix.unstack(-1).plot()

def curve_returns(index='IG', rolling=6):
    ## look at returns
    otr = on_the_run(index)
    df = index_returns(index=index, series=list(range(otr - rolling, otr + 1)),
                       tenor=['3yr', '5yr', '7yr', '10yr'])
    ## on-the-run returns
    df = df.reset_index().set_index(['date', 'series', 'tenor'])
    returns = df.price_return.dropna().unstack(-1).groupby(level='date').nth(-1)

    strategies_return = pd.DataFrame(
        {'5-10': 1.78 * returns['5yr'] - returns['10yr'],
         '7-10': 1.33 * returns['7yr'] - returns['10yr'],
         '3-5-10': -2 * returns['3yr'] + 3 * returns['5yr'] - returns['10yr'],
         '3-5': returns['5yr'] - 1.56 * returns['3yr'],
         '3-7': returns['7yr'] - 2.07 * returns['3yr']})
    strategies_return_monthly = (strategies_return.
                                 groupby(pd.Grouper(freq='M')).
                                 agg(lambda df: (1 + df).prod() - 1))

    def sharpe(df, period="daily"):
        if period == "daily":
            return df.mean() / df.std() * math.sqrt(252)
        else:
            return df.mean() / df.std() * math.sqrt(12)

    results = strategies_return.agg([sharpe, lambda df: df.nsmallest(10).mean()])
    sharpe_monthly = strategies_return_monthly.agg(sharpe, period="monthly")
    sharpe_monthly.name = 'Monthly Sharpe'
    results.index=['Sharpe', 'Mean Worst 10 Days DrawDown']
    return results.append(sharpe_monthly)

def cross_series_curve(index='IG', rolling=6):
    otr = on_the_run(index)
    df = index_returns(index= index, series=list(range(otr - rolling, otr + 1)),
                       tenor=['3yr', '5yr', '7yr', '10yr'])
    ## look cross series - 3y to 5y
    df = df.reset_index().set_index(['date', 'index', 'tenor', 'series'])
    returns1 = df.xs(['5yr', index], level = ['tenor','index']).price_return.unstack(-1)
    price_diff = pd.DataFrame()
    for ind in list(range(otr - 2, otr + 1)):
        price_diff[ind] = returns1[ind] - 1.6 * returns1[ind - 4]

    price_diff = price_diff.stack().groupby(level = 'date').nth(-1)
    monthly_returns_cross_series = (price_diff.
                                    groupby(pd.Grouper(freq='M')).
                                    agg(lambda df: (1 + df).prod() - 1))
    plt.plot(monthly_returns_cross_series)


def forward_loss(index='IG'):
    start_date = (pd.Timestamp.now() - pd.DateOffset(years=3)).date()

    df = pd.read_sql_query("SELECT date, index, series, tenor, duration, closespread, "\
                           "closespread*duration / 100 AS indexel " \
                           "FROM index_quotes WHERE index=%s AND date >= %s " \
                           "ORDER BY date DESC, series ASC, duration ASC",
                           serenitasdb, parse_dates=['date'], params=[index, start_date])
    df1 = pd.read_sql_query("SELECT index, series, tenor, maturity FROM index_maturity",
                            serenitasdb, parse_dates=['maturity'])

    df = df.merge(df1, on=['index','series','tenor'])
    df = df.set_index(['date','index', 'maturity']).dropna()
    df = df.groupby(level=['date','index', 'maturity']).nth(-1)
    # annual change, to take out some noise
    df['fwd_loss_rate'] = df.indexel.diff(2)/df.duration.diff(2)

def curve_model(tenor_1='5yr', tenor_2='10yr'):
    #OLS model
    df = ratio_within_series(param='closespread')
    df = pd.concat([df.duration[tenor_1], df.duration[tenor_2],
                    df.closespread[tenor_1],
                    df.closespread_ratio_to_5yr[tenor_2],
                    df.theta[tenor_1], df.theta[tenor_2]],
                   axis=1,
                   keys=['duration1', 'duration2', 'closespread',
                         'ratio', 'theta1', 'theta2'])
    df = np.log(df)
    ols_model = smf.ols('ratio ~ closespread + duration1 + theta1 + theta2',
                        data=df).fit()
    return df, ols_model

def curve_model_results(df, model):
    df = df.dropna()
    prstd_ols, df['down_2_stdev'], df['up_2_stdev'] = wls_prediction_std(model)
    #dr/dspread = exp(k) + spread_coeff * duration ^ dur_coeff * spread ^ (spread_coeff-1)
    cols = ['ratio', 'closespread', 'down_2_stdev', 'up_2_stdev']
    df[cols] = np.exp(df[cols])
    df['predicted'] = np.exp(model.predict())
    df[['predicted', 'down_2_stdev', 'up_2_stdev']]=\
        df[['predicted', 'down_2_stdev', 'up_2_stdev']].multiply(df['closespread'].values, axis=0)
    ax = df[['predicted', 'down_2_stdev', 'up_2_stdev']].reset_index(level='series', drop=True).plot()
    df['dr_dspread'] = np.exp(model.params[0]) * model.params[2] * df.duration1 ** model.params[1] * df.closespread ** (model.params[2] - 1)
    return df

def spread_fin_crisis(index='IG'):
    otr = on_the_run(index)
    ## look at spreads
    df = get_index_quotes(index, list(range(8, otr + 1)),
                          tenor=['3yr', '5yr', '7yr', '10yr'], years=20)
    spreads = df.groupby(level=['date', 'tenor']).nth(-1)['closespread'].unstack(-1)
    spreads_diff = spreads.diff(axis=1)
    to_plot = pd.DataFrame()
    to_plot['spread'] = spreads['5yr']
    to_plot['3 - 5 diff'] = spreads_diff['5yr']
    to_plot['5 - 10 diff'] = spreads_diff['7yr'] + spreads_diff['10yr']

    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax2 = ax.twinx() # Create another axes that shares the same x-axis as ax

    width = 0.4
    to_plot['spread'].plot(color='red', ax=ax)
    to_plot['5 - 10 diff'].plot(color='blue', ax=ax2)
    to_plot['3 - 5 diff'].plot(color='green', ax=ax2)
    plt.legend(bbox_to_anchor=(.5, -.1), ncol = 2)

    plt.show()

def forward_spread(index='IG', series=None, tenors=['3yr', '5yr', '7yr', '10yr']):

    if series is None:
        series = on_the_run(index = index)
    b_index = MarkitBasketIndex(index, series, tenors)
    b_index.tweak()

    f_spread = []
    date_range = pd.bdate_range(pd.datetime.today(), max(b_index.maturities), freq='M')
    for d in date_range.date:
        b_index.trade_date = d
        f_spread.append(b_index.spread())
    return pd.concat(f_spread, keys=date_range).unstack(-1)

def spot_forward(index='IG', series=None, tenors=['3yr', '5yr', '7yr', '10yr']):

    if series is None:
        series = on_the_run(index)
    b_index = MarkitBasketIndex(index, series, tenors)
    b_index.tweak()

    spreads_current = b_index.spread()
    spreads_current.name = 'current'
    spreads_1yr = pd.Series([b_index.spread(m - relativedelta(years=1), b_index.coupon(m)) \
                             for m in b_index.maturities], index=tenors)
    spreads_1yr.name = '1yr'
    df = pd.concat([spreads_current, spreads_1yr], axis=1)
    maturity_1yr = roll_date(b_index.index_desc.issue_date[0], 1)
    df_0 = pd.DataFrame({'current':[0., b_index.spread(maturity_1yr,
                                                       0.01 if index == "IG" else 0.05)],
                         '1yr': [0., 0.]}, index=['0yr', '1yr'])
    df_0.index.name = 'tenor'
    df = df_0.append(df)
    df['maturity'] = [b_index.trade_date, maturity_1yr] + b_index.maturities
    return df.reset_index().set_index('maturity')