import pandas as pd
import argparse
from analytics import Index, Swaption
from db import dbconn
from joblib import Parallel, delayed
from pickle import loads, dumps

serenitasdb = dbconn('serenitasdb')

def get_data(index, series, date = "'-infinity'::date"):
    df = pd.read_sql("SELECT * from swaption_ref_quotes JOIN swaption_quotes " \
                     "USING (quotedate, index, series, expiry) WHERE index=%s and series=%s " \
                     "and quotedate >=%s ORDER BY quotedate",
                     "postgresql://serenitas_user@debian/serenitasdb",
                     params = (index, series, date), parse_dates = ['quotedate', 'expiry'])
    df.loc[(df.quote_source == "GS") & (df['index'] =="HY"),
           ["pay_bid", "pay_offer", "rec_bid", "rec_offer"]] *=100
    return df


def calib(d, option, expiry, index, series):
    option.strike = d['strike']
    option.ref = d['ref']
    r = []
    for pv_type in ['pv', 'pv_black']:
        for option_type in ['pay', 'rec']:
            mid = (d['{}_bid'.format(option_type)] + d['{}_offer'.format(option_type)])/2 * 1e-4
            option.option_type = 'payer' if option_type == 'pay' else 'receiver'
            try:
                setattr(option, pv_type, mid)
            except ValueError as e:
                r.append(None)
                print(e)
            else:
                r.append(option.sigma)
    return [d['quotedate'], index, series, expiry, d['strike']] + r

def calibrate(index_type, series, date):
    sql_str = ("INSERT INTO swaption_calib VALUES({}) ON CONFLICT DO NOTHING".
               format(",".join(["%s"] * 9)))
    data = get_data(index_type, series, date)
    index = Index.from_name(index_type, series, "5yr")

    for k, v in data.groupby([data['quotedate'].dt.date, 'expiry']):
        trade_date, expiry = k
        index.trade_date = trade_date
        option = Swaption(index, expiry.date(), 100,
                          strike_is_price = index_type == "HY")
        r = Parallel(n_jobs=4)(
            delayed(calib)(d, option, expiry.date(), index_type, series) for d in
            v[['ref', 'quotedate', 'strike', 'pay_bid', 'pay_offer', 'rec_bid', 'rec_offer']].
            to_dict(orient = 'records'))
        with serenitasdb.cursor() as c:
            c.executemany(sql_str, r)
            serenitasdb.commit()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--index', required=True, type = lambda s: s.upper())
    parser.add_argument('--series', required=True, type=int)
    parser.add_argument('--date', required = False, default = "'-infinity'::date")
    args = parser.parse_args()
    calibrate(args.index, args.series, args.date)