import datetime
from bbg_helpers import init_bbg_session, retrieve_data, BBG_IP
from analytics.utils import tenor_t
import pandas as pd

spreads = [-200, -100, -50, -25, 25, 50, 100, 200]
spread_letter = {200: 'D', 100: 'C', 75: 'G', 50: 'B', 25: 'A'}


def ticker(expiry, tenor, spread, vol_type, source):
    if spread < 0:
        spread_mapping = "R" + spread_letter[abs(spread)]
    else:
        spread_mapping = "P" + spread_letter[abs(spread)]
    if vol_type == "V":
        return f"US{spread_mapping}{expiry:0>2}{tenor:0>2} {source} Curncy"
    elif vol_type == "N":
        return f"USS{spread_mapping}{expiry}{tenor:0>2} {source} Curncy"

# Normal vol
# sources GFIS, CMPN
expiry = ["C", "F", 1, 2, 3, 4, 5, 7, 10, 20, 30]
tenor = [1, 2, 5, 10, 20, 30]
tickers = []
for s in spreads:
    for e in expiry:
        for t in tenor:
            tickers.append(ticker(e, t, s, "V", "GFIS"))

with init_bbg_session(BBG_IP) as session:
    data = retrieve_data(session, tickers, ['PX_LAST'],
                         start_date=datetime.date(2018, 1, 1))

df = pd.concat(data, axis=1)
df.columns = df.columns.droplevel(1)
col = df.columns.to_series().str.extract("US([RP])([DCGBA])(\d[0-9ACF])(\d{2})")
pos_dict = {v: k for k, v in spread_letter.items()}
neg_dict = {v: -k for k, v in spread_letter.items()}
col.loc[col[0] == 'P', 1] = col.loc[col[0] == 'P', 1].map(pos_dict)
col.loc[col[0] == 'R', 1] = col.loc[col[0] == 'R', 1].map(neg_dict)

def to_tenor(s):
    d = {'A': '1m', 'C': '3m', 'F': '6m'}
    return d.get(s, s + 'yr')

col[2] = col[2].str.lstrip("0").map(to_tenor).astype(tenor_t)
col[3] = col[3].str.lstrip("0").map(to_tenor).astype(tenor_t)
df.columns = pd.MultiIndex.from_arrays([col[1], col[2], col[3]],
                                       names=['spread', 'expiry', 'tenor'])