from db import dbconn
import datetime
from io import BytesIO
import os
import pandas as pd
import requests
import xml.etree.ElementTree as ET
import zipfile
from yieldcurve import YC, ql_to_jp
from pickle import dumps

def downloadMarkitIRData(download_date=datetime.date.today(),
                         currency="USD"):
    conn = dbconn("serenitasdb")
    ## T+1 rates are published in the evening
    effective_date = download_date + datetime.timedelta(days=1)
    basedir = os.path.join(os.environ['DATA_DIR'], "Yield Curves")
    filename = "InterestRates_{0}_{1:%Y%m%d}".format(currency, effective_date)
    if not os.path.exists(os.path.join(basedir, filename + '.xml')):
        r = requests.get('http://www.markit.com/news/{0}.zip'.format(filename))
        if "zip" in r.headers['content-type']:
            with zipfile.ZipFile(BytesIO(r.content)) as z:
                z.extractall(path = os.path.join(os.environ['DATA_DIR'], "Yield Curves"))
        else:
            return downloadMarkitIRData(download_date - datetime.timedelta(days=1))

    tree = ET.parse(os.path.join(os.environ['DATA_DIR'], "Yield Curves", filename + '.xml'))
    deposits = zip([e.text for e in tree.findall('./deposits/*/tenor')],
                   [float(e.text) for e in tree.findall('./deposits/*/parrate')])
    swaps = zip([e.text for e in tree.findall('./swaps/*/tenor')],
                [float(e.text) for e in tree.findall('./swaps/*/parrate')])
    effectiveasof = tree.find('./effectiveasof').text
    MarkitData = {'deposits': list(deposits),
                  'swaps': list(swaps),
                  'effectiveasof': pd.Timestamp(effectiveasof).date()}
    ql_yc = YC(currency=currency, MarkitData=MarkitData)
    jp_yc = ql_to_jp(ql_yc)
    sql_str = "INSERT INTO {}_curves VALUES(%s, %s) ON CONFLICT DO NOTHING".format(currency)
    with conn.cursor() as c:
        c.execute(sql_str, (MarkitData['effectiveasof'], jp_yc.__getstate__()))
    instruments = MarkitData['deposits'] + MarkitData['swaps']
    names = ",".join(['"{}"'.format(r[0]) for r in instruments])
    values = ",".join(["%s"] * (len(instruments) + 1)) # +1 for effective_date
    insert_str = ("INSERT INTO {0}_rates(effective_date, {1}) VALUES({2}) ON CONFLICT DO NOTHING".
                  format(currency, names, values))
    with conn.cursor() as c:
        c.execute(insert_str, [MarkitData['effectiveasof']] +[r[1] for r in instruments])
    conn.commit()
    conn.close()