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, calendar_from_name
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")
    filename = "InterestRates_USD_{%Y%m%d}.xml".format(date)
    if not os.path.exists(os.path.join(basedir, filename)):
        r = requests.get('http://www.markit.com/news/InterestRates_USD_{%Y%m%d}.zip'.format(date))
        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-{%Y-%m-%d}.csv".format(date)
    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(), 0).rate for d in sched]
    f6 = [ts.forward_rate(d, d+p6, Actual360(), 0).rate for d in sched]
    f2 = [ts.forward_rate(d, d+p2, Actual360(), 0).rate for d in sched]

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