from common import root
import os
import requests, zipfile
from io import BytesIO
import xml.etree.ElementTree as ET
import datetime
from quantlib.time.api import Calendar, Period, Days, Schedule, today, Actual360
from quantlib.time import imm
from quantlib.util.converter import qldate_to_pydate, pydate_to_qldate
from quantlib.market.market import libor_market, next_imm_date
import numpy as np
import matplotlib.pyplot as plt

def getMarkitIRData(date = datetime.date.today()):
    basedir = os.path.join(root, "data", "Yield Curves")
    datestr = date.strftime("%Y%m%d")
    filename = "InterestRates_USD_{0}.xml".format(datestr)
    if not os.path.exists(os.path.join(basedir, filename)):
        r = requests.get('http://www.markit.com/news/InterestRates_USD_{0}.zip'.format(datestr))
        if "x-zip" in r.headers['content-type']:
            with zipfile.ZipFile(BytesIO(r.content)) as z:
                z.extractall(path = os.path.join(root, "data", "Yield Curves"))
        else:
            return getMarkitIRData(date-datetime.timedelta(days=1))

    tree = ET.parse(os.path.join(root, "data", "Yield Curves", filename))
    deposits = {tenor: rate for tenor, rate in \
                zip([e.text for e in tree.findall('./deposits/*/tenor')],
                    [float(e.text) for e in tree.findall('./deposits/*/parrate')])}
    swaps = {tenor: rate for tenor, rate in \
             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': deposits,
                  'swaps': swaps,
                  'effectiveasof': datetime.datetime.strptime(effectiveasof, "%Y-%m-%d").date()}
    return MarkitData

def get_futures_data(date = datetime.date.today()):
    futures_file = os.path.join(root, "data", "Yield Curves",
                                "futures-{0}.csv".format(date.strftime("%Y-%m-%d")))
    with open(futures_file) as fh:
        quotes = [float(line.split(",")[1]) for line in fh]
    return quotes

def YC(date = datetime.date.today(), MarkitData=None, futures = None):
    if not MarkitData:
        MarkitData = getMarkitIRData(date)
    if not futures:
        futures = get_futures_data(date)
    m = libor_market('USD(NY)')
    cal = Calendar.from_name('GBR')
    # m.settle_date is not available until we set_quotes, so we compute it again
    # need a better way to do this
    settle_date = cal.advance(pydate_to_qldate(date), 2, Days)
    quotes = [('ED',i+1, v) for i, v in enumerate(futures)]
    if next_imm_date(date, 9) == settle_date + Period('2Yr'):
        quotes.pop(7)
    quotes += [('SWAP', k, v) for k, v in MarkitData['swaps'].items()]
    m.set_quotes(date, quotes)
    ts = m.bootstrap_term_structure()
    return ts

if __name__=="__main__":
    date = datetime.date(2014, 4, 29)
    ts = YC(date)
    cal = Calendar.from_name('USA')
    p1 = Period('1Mo')
    p2 = Period('2Mo')
    p3 = Period('3Mo')
    p6 = Period('6Mo')
    p12 = Period('12Mo')
    sched = Schedule(ts.reference_date, ts.reference_date+Period('5Yr'), Period('3Mo'), cal)
    days = [qldate_to_pydate(day) for day in sched]
    f3 = [ts.forward_rate(d, d+p3, Actual360()).rate for d in sched]
    f6 = [ts.forward_rate(d, d+p6, Actual360()).rate for d in sched]
    f2 = [ts.forward_rate(d, d+p2, Actual360()).rate for d in sched]

    plt.plot(days, f2, days, f3, days, f6)
    plt.show()