import pandas as pd
import argparse
from analytics import Index, Swaption
import datetime
from db import dbconn
from joblib import Parallel, delayed
from pickle import loads, dumps
from psycopg2.extensions import register_adapter

serenitasdb = dbconn('serenitasdb')

def get_data(index, series, date = datetime.date.min):
    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, nproc):
    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=nproc)(
            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 = datetime.date.min)
    parser.add_argument('--nproc', required = False, type=int, default = 4)
    args = parser.parse_args()
    calibrate(args.index, args.series, args.date, args.nproc)