from db import dbconn
import datetime
from io import BytesIO
import lz4
import os
import pandas as pd
from psycopg2 import sql
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 = f"InterestRates_{currency}_{effective_date:%Y%m%d}"
    if not os.path.exists(os.path.join(basedir, filename + '.xml')):
        r = requests.post(f'http://www.markit.com/news/{filename}.zip')
        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, evaluation_date=MarkitData['effectiveasof'])
    jp_yc = ql_to_jp(ql_yc)
    sql_str = f"INSERT INTO {currency}_curves VALUES(%s, %s) ON CONFLICT DO NOTHING"
    with conn.cursor() as c:
        c.execute(sql_str, (MarkitData['effectiveasof'],
                            lz4.block.compress(jp_yc.__getstate__())))
    instruments = MarkitData['deposits'] + MarkitData['swaps']
    names = sql.SQL(", ").join([sql.Identifier(r[0]) for r in instruments])
    values = sql.SQL(", ").join(sql.Placeholder() * (len(instruments) + 1)) # +1 for effective_date
    insert_str = (sql.SQL(f"INSERT INTO {currency}_rates(effective_date, {{}}) VALUES({{}}) " \
                          "ON CONFLICT DO NOTHING").
                  format(names, values))
    with conn.cursor() as c:
        c.execute(insert_str, [MarkitData['effectiveasof']] +[r[1] for r in instruments])
    conn.commit()
    conn.close()