from blist import sorteddict
from collections import namedtuple
from contextlib import closing
from itertools import islice
import datetime
import os
import pandas as pd
from quantlib.settings import Settings
from quantlib.time.api import (
    WeekendsOnly,
    Japan,
    Date,
    Period,
    Days,
    Schedule,
    Annual,
    Semiannual,
    today,
    Actual360,
    Months,
    Years,
    ModifiedFollowing,
    Thirty360,
    Actual365Fixed,
)
from quantlib.currency.api import USDCurrency, EURCurrency, JPYCurrency
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
from quantlib.quotes import SimpleQuote
from utils.db import dbconn, serenitas_engine
from pyisda.curve import YieldCurve
from pyisda.date import BadDay
import warnings


def load_curves(currency="USD", date=None):
    """load the prebuilt curves from the database"""
    sql_str = f"SELECT * FROM {currency}_curves"
    if date:
        sql_str += " WHERE effective_date=%s"

    with closing(dbconn("serenitasdb")) as conn:
        with conn.cursor() as c:
            if date:
                c.execute(sql_str, (date,))
                if c:
                    _, curve = c.fetchone()
                    return YieldCurve.from_bytes(curve)
            else:
                c.execute(sql_str)
                return sorteddict([(d, YieldCurve.from_bytes(curve)) for d, curve in c])


def get_curve(effective_date, currency="USD"):
    if f"_{currency}_curves" in globals():
        curves = globals()[f"_{currency}_curves"]
    else:
        curves = globals()[f"_{currency}_curves"] = load_curves(currency)
    if isinstance(effective_date, datetime.datetime):
        effective_date = effective_date.date()
    if effective_date > curves.keys()[-1]:
        last_curve = curves[curves.keys()[-1]]
        return last_curve
    if effective_date in curves:
        return curves[effective_date]
    else:
        warnings.warn(
            f"cache miss for {currency} curve on {effective_date}", RuntimeWarning
        )
        ql_yc = YC(currency=currency, evaluation_date=effective_date)
        jp_yc = ql_to_jp(ql_yc)
        curves[effective_date] = jp_yc
        return jp_yc


def getMarkitIRData(effective_date=datetime.date.today(), currency="USD"):
    conn = dbconn("serenitasdb")
    sql_str = (
        "SELECT * FROM {}_rates WHERE  effective_date <= %s "
        "ORDER BY effective_date DESC LIMIT 1".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, rate) for t, rate in zip(col_names[1:7], r[1:7]) if rate is not None
        ],
        "swaps": [
            (t, rate) for t, rate in zip(col_names[7:], r[7:]) if rate is not None
        ],
    }
    return MarkitData


def get_futures_data(date=datetime.date.today()):
    futures_file = os.path.join(
        os.environ["DATA_DIR"], "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 get_curve_params(currency):
    if currency == "USD":
        currency = USDCurrency()
        calendar = WeekendsOnly()
        fixed_dc = Thirty360()
        floating_dc = Actual360()
        mm_dc = Actual360()
        floating_freq = Period(3, Months)
        fixed_freq = Semiannual
    elif currency == "EUR":
        currency = EURCurrency()
        calendar = WeekendsOnly()
        fixed_dc = Thirty360()
        floating_dc = Actual360()
        mm_dc = Actual360()
        floating_freq = Period(6, Months)
        fixed_freq = Annual
    elif currency == "JPY":
        currency = JPYCurrency()
        calendar = Japan()
        fixed_dc = Actual365Fixed()
        floating_dc = Actual360()
        mm_dc = Actual360()
        floating_freq = Period(6, Months)
        fixed_freq = Semiannual
    CurveParams = namedtuple(
        "CurveParam",
        "currency, calendar, fixed_dc, floating_dc, "
        "mm_dc, floating_freq, fixed_freq",
    )
    return CurveParams(
        currency, calendar, fixed_dc, floating_dc, mm_dc, floating_freq, fixed_freq
    )


def rate_helpers(currency="USD", MarkitData=None, evaluation_date=None):
    """Utility 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 evaluation_date is None:
        evaluation_date = pydate_from_qldate(settings.evaluation_date)
    if isinstance(evaluation_date, pd.Timestamp):
        evaluation_date = evaluation_date.date()
    if not MarkitData:
        MarkitData = getMarkitIRData(evaluation_date, currency)
        if MarkitData["effectiveasof"] != evaluation_date:
            warnings.warn(
                "Yield curve effective date: {0} doesn't "
                "match the evaluation date: {1}".format(
                    MarkitData["effectiveasof"], evaluation_date
                ),
                RuntimeWarning,
            )
    settings.evaluation_date = Date.from_datetime(MarkitData["effectiveasof"])
    params = get_curve_params(currency)
    isda_ibor = IborIndex(
        "IsdaIbor",
        params.floating_freq,
        2,
        params.currency,
        params.calendar,
        ModifiedFollowing,
        False,
        params.floating_dc,
    )
    # we use SimpleQuotes, rather than just float to make it updateable
    deps = [
        DepositRateHelper(
            SimpleQuote(q),
            Period(t),
            2,
            params.calendar,
            ModifiedFollowing,
            False,
            params.mm_dc,
        )
        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),
            params.calendar,
            params.fixed_freq,
            ModifiedFollowing,
            params.fixed_dc,
            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 = [d for d in yc.dates if d >= 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,
    evaluation_date=None,
    fixed=False,
    extrapolation=False,
):
    calendar = WeekendsOnly()
    settings = Settings()
    if evaluation_date:
        settings.evaluation_date = Date.from_datetime(evaluation_date)
    if helpers is None:  # might roll back the evaluation date
        helpers = rate_helpers(currency, MarkitData, evaluation_date)

    if fixed:
        _yc = PiecewiseYieldCurve.from_reference_date(
            BootstrapTrait.Discount,
            Interpolator.LogLinear,
            settings.evaluation_date,
            helpers,
            Actual365Fixed(),
        )
    else:
        _yc = PiecewiseYieldCurve(
            BootstrapTrait.Discount,
            Interpolator.LogLinear,
            0,
            calendar,
            helpers,
            Actual365Fixed(),
        )
    if extrapolation:
        _yc.extrapolation = True
    return _yc


def jpYC(effective_date, currency="USD", MarkitData=None):
    if MarkitData is None:
        markit_data = getMarkitIRData(effective_date, currency)
    periods, rates = zip(*markit_data["deposits"])
    periods_swaps, rates_swaps = zip(*markit_data["swaps"])
    types = "M" * len(periods) + "S" * len(periods_swaps)
    rates = np.array(rates + rates_swaps)
    periods = list(periods + periods_swaps)
    if currency == "USD":
        fixed_period = "6M"
        float_period = "3M"
    elif currency == "EUR":
        fixed_period = "12M"
        float_period = "6M"
    return YieldCurve(
        effective_date,
        types,
        periods,
        rates,
        "ACT/360",
        fixed_period,
        float_period,
        "30/360",
        "ACT/360",
        BadDay.MODIFIED,
    )


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 RuntimeError("QuantLib curve needs to use Discount trait")


def build_curves(currency="USD"):
    settings = Settings()
    params = get_curve_params(currency)
    isda_ibor = IborIndex(
        "IsdaIbor",
        params.floating_freq,
        2,
        params.currency,
        params.calendar,
        ModifiedFollowing,
        False,
        params.floating_dc,
    )
    rates = pd.read_sql_table(
        f"{currency.lower()}_rates", serenitas_engine, index_col="effective_date"
    )
    quotes = [SimpleQuote() for c in rates.columns]
    gen = zip(quotes, rates.columns)
    deps = [
        DepositRateHelper(
            q, Period(t), 2, params.calendar, ModifiedFollowing, False, params.mm_dc
        )
        for q, t in islice(gen, 6)
    ]
    swaps = [
        SwapRateHelper.from_tenor(
            q,
            Period(t),
            params.calendar,
            params.fixed_freq,
            ModifiedFollowing,
            params.fixed_dc,
            isda_ibor,
        )
        for q, t in gen
    ]
    sql_str = f"INSERT INTO {currency}_curves VALUES(%s, %s) ON CONFLICT DO NOTHING"
    conn = serenitas_engine.raw_connection()
    for effective_date, curve_data in rates.iterrows():
        print(effective_date)
        settings.evaluation_date = Date.from_datetime(effective_date)
        for q, val in zip(quotes, curve_data):
            q.value = val
        valid_deps = [d for d in deps if not np.isnan(d.quote)]
        valid_swaps = [s for s in swaps if not np.isnan(s.quote)]
        ql_yc = PiecewiseYieldCurve(
            BootstrapTrait.Discount,
            Interpolator.LogLinear,
            0,
            calendar,
            valid_deps + valid_swaps,
            Actual365Fixed(),
        )
        jp_yc = ql_to_jp(ql_yc)
        with conn.cursor() as c:
            c.execute(
                sql_str, (effective_date, lz4.block.compress(jp_yc.__getstate__()))
            )
    conn.commit()
    conn.close()


if __name__ == "__main__":
    # evaluation_date = Date(29, 4, 2014)
    Settings.instance().evaluation_date = today()

    import matplotlib.pyplot as plt
    from quantlib.time.api import calendar_from_name
    from pandas.plotting import register_matplotlib_converters

    register_matplotlib_converters()
    helpers = rate_helpers("USD")
    ts = YC(helpers)
    cal = calendar_from_name("USA")
    p1 = Period("1M")
    p2 = Period("2M")
    p3 = Period("3M")
    p6 = Period("6M")
    p12 = Period("12M")
    sched = Schedule.from_rule(
        ts.reference_date, ts.reference_date + Period("5Y"), Period("3M"), 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)