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
from quantlib.quotes import SimpleQuote

def getMarkitIRData(date = datetime.date.today(), currency="USD"):
    basedir = os.path.join(root, "data", "Yield Curves")
    filename = "InterestRates_{0}_{1:%Y%m%d}".format(currency, 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 "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 + '.xml'))
    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:%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(), currency="USD", MarkitData=None, futures = None):
    if not MarkitData:
        MarkitData = getMarkitIRData(date, currency)
    # if not futures:
    #     futures = get_futures_data(date)
    if currency == "USD":
        m = libor_market('USD(NY)', calendar='WO')
    else if currency == "EUR":
        m = libor_market('EUR:1Y', calendar='WO')
        m.calendar = "WeekendsOnly"
        m.fixed_leg_convention = 'ModifiedFollowing'
    # m.settle_date is not available until we set_quotes, so we compute it again
    # need a better way to do this
    # 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 = [('DEP', t, SimpleQuote(v)) for t, v in MarkitData['deposits'].items()]
    quotes += [('SWAP', t, SimpleQuote(v)) for t, 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(d) for d 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()