#we want a function that rebook trades
from sqlalchemy import create_engine
import pandas as pd
from pandas.tseries.offsets import BDay
import pdb

def get_outstanding_positions(engine):
    sqlstr = """WITH current_portfolio AS (
SELECT security_id, maturity, attach, detach, sum(notional * (CASE WHEN protection='Buyer' THEN -1 ELSE 1 END))
AS ntl FROM cds
GROUP BY  security_id, maturity, attach, detach)
SELECT distinct security_id, lastdate FROM current_portfolio JOIN index_desc
ON current_portfolio.security_id=index_desc.redindexcode
AND index_desc.maturity=current_portfolio.maturity WHERE lastdate!='infinity' AND
ntl !=0 ORDER BY lastdate"""
    df = pd.read_sql_query(sqlstr, engine)
    return df

def next_redcode(engine, secid):
    sqlstr = """SELECT DISTINCT redindexcode FROM index_desc JOIN
    (SELECT DISTINCT index, series, version+1 AS version FROM index_desc WHERE redindexcode=%s) a
    USING (index, series, version)"""
    conn = engine.raw_connection()
    with conn.cursor() as c:
        c.execute(sqlstr, (secid,))
        nextredcode, = c.fetchone()
    conn.commit()
    conn.close()
    return nextredcode

def rebook(engine, secid, lastdate, already_rebooked):
    trades = pd.read_sql("SELECT * from cds WHERE security_id = %s", engine, params=(secid,))
    next_secid = next_redcode(engine, secid)
    for i in range(trades.shape[0]):
        dealid = trades.at[i,'dealid']
        if (secid, dealid) in already_rebooked:
            continue
        trade = trades.iloc[[i]]
        trade = trade.append(trade, ignore_index=True)
        del trade['id']
        del trade['lastupdate']
        trade['trade_date'] = lastdate + BDay(1)
        trade['cp_code'] = '999'
        trade.set_value(0, 'protection', 'Buyer' if trade.at[0,'protection'] == 'Seller' else 'Seller')
        trade.set_value(1, 'security_id', next_secid)
        trade.to_sql('cds', engine, if_exists='append', index=False)
    return (secid, dealid)

def rebook_everything(engine):
    torebook = get_outstanding_positions(engine)
    already_rebooked = set([])
    while not torebook.empty:
        pdb.set_trace()
        secid , lastdate = torebook.iloc[0].tolist()
        already_rebooked.add(rebook(engine, secid, lastdate, already_rebooked))
        torebook = get_outstanding_positions(engine)

if __name__=="__main__":
    engine = create_engine("postgresql://dawn_user@debian/dawndb")
    rebook_everything(engine)