from common import root
import os
import datetime
from quantlib.settings import Settings
from quantlib.time.api import (WeekendsOnly, Date, Period, Days, Schedule, Annual,
                               Semiannual, today, Actual360, Months, Years,
                               ModifiedFollowing, Thirty360, Actual365Fixed,
                               calendar_from_name)
from quantlib.currency.api import USDCurrency, EURCurrency
from quantlib.indexes.ibor_index import IborIndex
from quantlib.termstructures.yields.api import (
    PiecewiseYieldCurve, DepositRateHelper, SwapRateHelper, BootstrapTrait, Interpolator)
from quantlib.time.date import pydate_from_qldate

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from quantlib.quotes import SimpleQuote
from db import dbconn
from pyisda.curve import YieldCurve

def getMarkitIRData(effective_date = datetime.date.today(),
                    currency = "USD"):
    conn = dbconn("serenitasdb")
    sql_str = "SELECT * FROM {}_rates WHERE  effective_date = %s".format(currency)
    with conn.cursor() as c:
        c.execute(sql_str, (effective_date,))
        col_names = [col[0] for col in c.description]
        r = c.fetchone()
    MarkitData = {'effectiveasof': r[0],
                  'deposits': [(t, r[i]) for i, t in \
                               enumerate(col_names[1:7], 1) if r[i] is not None],
                  'swaps': [(t, r[i]) for i, t in enumerate(col_names[7:], 7)]}
    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 rate_helpers(currency="USD", MarkitData=None):
    """Util function to build a list of RateHelpers

    Parameters
    ----------
    currency : str, optional
       One of `USD`, `EUR` at the moment, defaults to `USD`
    MarkitData : dict, optional
       MarkitData for the current evaluation_date

    Returns
    -------
    helpers : list
       List of QuantLib RateHelpers
    """
    settings = Settings()
    if not MarkitData:
        MarkitData = getMarkitIRData(pydate_from_qldate(settings.evaluation_date), currency)
        if MarkitData['effectiveasof'] != pydate_from_qldate(settings.evaluation_date):
            raise RuntimeError("Yield curve effective date: {0} doesn't " \
                               "match the evaluation date: {1}".format(
                                   MarkitData['effectiveasof'],
                                   pydate_from_qldate(settings.evaluation_date)))
    calendar = WeekendsOnly()
    if currency == "USD":
        isda_ibor = IborIndex("IsdaIbor", Period(3, Months), 2, USDCurrency(), calendar,
                              ModifiedFollowing, False, Actual360())
        fix_freq = Semiannual
    elif currency == "EUR":
        isda_ibor = IborIndex("IsdaIbor", Period(6, Months), 2, EURCurrency(), calendar,
                              ModifiedFollowing, False, Actual360())
        fix_freq = Annual
    # we use SimpleQuotes, rather than just float to make it updateable
    deps = [DepositRateHelper(SimpleQuote(q), Period(t), 2, calendar, ModifiedFollowing, False, Actual360())
            for t, q in MarkitData['deposits']]
    # this matches with bloomberg, but according to Markit, maturity should be unadjusted
    swaps = [SwapRateHelper.from_tenor(SimpleQuote(q), Period(t), calendar, fix_freq, ModifiedFollowing,
                                       Thirty360(), isda_ibor) for t, q in MarkitData['swaps']]
    return deps + swaps

def get_dates(date, currency="USD"):
    """computes the list of curve dates on a given date"""
    if currency == "USD":
        month_periods = [1, 2, 3, 6, 12]
    year_periods = [2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 15, 20, 25, 30]
    calendar = WeekendsOnly()
    settle_date = calendar.advance(Date.from_datetime(date), 2, 0)
    deposit_dates = [calendar.advance(settle_date, period = Period(m, Months),
                                      convention=ModifiedFollowing) \
                     for m in month_periods]
    swap_dates = [calendar.advance(settle_date, period = Period(y, Years),
                                   convention=ModifiedFollowing) \
                     for y in year_periods]
    dates = deposit_dates + swap_dates
    return [pydate_from_qldate(d) for d in dates]

def roll_yc(yc, forward_date):
    """returns the expected forward yield cuve on a forward_date"""
    dates = get_dates(forward_date)
    dfs = np.array([yc.discount_factor(d, forward_date) for d in dates])
    return YieldCurve.from_discount_factors(forward_date, dates, dfs, 'ACT/365F')

def YC(helpers = None, currency="USD", MarkitData=None):
    if helpers is None:
        helpers = rate_helpers(currency, MarkitData)
    calendar = WeekendsOnly()
    return PiecewiseYieldCurve(BootstrapTrait.Discount, Interpolator.LogLinear,
                               0, calendar, helpers, Actual365Fixed())
def ql_to_jp(ql_yc):
    """convert a QuantLib yield curve to a JP's one"""
    if ql_yc._trait == BootstrapTrait.Discount:
        dfs = np.array(ql_yc.data[1:])
        dates = [pydate_from_qldate(d) for d in ql_yc.dates[1:]]
        trade_date = pydate_from_qldate(ql_yc.dates[0])
        return YieldCurve.from_discount_factors(trade_date, dates, dfs, 'ACT/365F')
    else:
        raise RuntimeErrror('QuantLib curve needs to use Discount trait')

if __name__=="__main__":
    #evaluation_date = Date(29, 4, 2014)
    Settings.instance().evaluation_date = today()
    ts = YC()
    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 = [pydate_from_qldate(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)