from bbg_helpers import init_bbg_session, retrieve_data, BBG_IP
import pandas as pd
import datetime
from db import dbengine, dbconn
from dateutil.relativedelta import relativedelta

def get_universe():
    engine = dbengine('corelogic')
    r = engine.execute("SELECT cusip from map_cusip")
    return set(c[0][:8] for c in r.fetchall() if c[0])

def get_bbg_data(cusips, universe):
    end_date = datetime.datetime.utcnow()
    start_date = end_date - relativedelta(years=1)
    with init_bbg_session(BBG_IP) as session:
        for cusip in cusips:
            if cusip in universe:
                security = "{}@TRAC Mtge".format(cusip)
                yield (cusip, retrieve_data(session, security,
                                            start_date=start_date,
                                            end_date=end_date))
def insert_data(dfs, conn):
    sql_str = "INSERT INTO trace_trades(cusip, condition_code, size, time, price)" \
              " VALUES(%s, %s, %s, %s, %s) ON CONFLICT DO NOTHING"
    for cusip, df in dfs:
        if 'type' in df:
            del df['type']
        if 'conditionCodes' not in df:
            df['conditionCodes'] = None
        df = df.sort_index(axis=1)
        cusip = cusip + calc_check_digit(cusip)
        to_insert = [(cusip,) + r for r in df.itertuples(index=False)]
        with conn.cursor() as c:
            c.executemany(sql_str, to_insert)
        conn.commit()

def get_cusips():
    with open("/home/share/CorpCDOs/data/TRACE/ABS.csv") as fh:
        for line in fh:
            yield line.split()[0]

_alphabet = '0123456789ABCDEFGH JKLMN PQRSTUVWXYZ*@#'

def calc_check_digit(number):
    """Calculate the check digits for the number."""
    # convert to numeric first, then sum individual digits
    number = ''.join(
        str((1, 2)[i % 2] * _alphabet.index(n)) for i, n in enumerate(number))
    return str((10 - sum(int(n) for n in number)) % 10)

if __name__ == "__main__":
    universe = get_universe()
    cusips = get_cusips()
    dfs = get_bbg_data(cusips, universe)
    conn = dbconn('serenitasdb')
    insert_data(dfs, conn)
    conn.close()