from analytics.basket_index import MarkitBasketIndex
from pyisda.legs import FeeLeg, ContingentLeg
from pyisda.logging import enable_logging

import pandas as pd

from db import dbconn


def all_curves_pv(curves, today_date, jp_yc, start_date, step_in_date, value_date, maturities):
    r = {}
    for d in maturities:
        tenor = {}
        coupon_leg = FeeLeg(start_date, d, True, 1., 1.)
        default_leg = ContingentLeg(start_date, d, True)
        accrued = coupon_leg.accrued(step_in_date)
        tickers = []
        data = []
        for sc in curves:
            coupon_leg_pv = coupon_leg.pv(today_date, step_in_date, value_date, jp_yc, sc, False)
            default_leg_pv = default_leg.pv(today_date, step_in_date, value_date,
                                            jp_yc, sc, 0.4)
            tickers.append(sc.ticker)
            data.append((coupon_leg_pv-accrued, default_leg_pv))
        r[pd.Timestamp(d)] = pd.DataFrame.from_records(data,
                                                       index=tickers,
                                                       columns=['duration', 'protection_pv'])
    return pd.concat(r, axis=1).swaplevel(axis=1).sort_index(axis=1, level=0)


def calibrate_portfolio(index_type, series, tenors=['3yr', '5yr', '7yr', '10yr'],
                        start_date=None):
    try:
        index = MarkitBasketIndex(index_type, series, tenors)
    except ValueError:
        return
    if start_date:
        index.index_quotes = index.index_quotes[start_date:]
    for value_date, v in index.index_quotes.groupby('date')['id']:
        index.value_date = value_date
        index.tweak()
        df = pd.concat([index.theta(),
                        index.duration(),
                        pd.Series(index.tweaks, index=tenors, name='tweak')], axis=1)
        for (_, t), id in v.items():
            yield (id, df.loc[t])


if __name__ == "__main__":
    enable_logging()
    import argparse
    import logging
    import os
    parser = argparse.ArgumentParser()
    parser.add_argument('index', help="index type (IG, HY, EU or XO)")
    parser.add_argument('series', help="series", type=int)
    parser.add_argument('--latest', required=False, action="store_true")
    args = parser.parse_args()
    index, series = args.index, args.series
    conn = dbconn('serenitasdb')

    if args.latest:
        with conn.cursor() as c:
            c.execute("SELECT max(date) FROM index_quotes_pre "
                      "RIGHT JOIN index_risk2 USING (id) "
                      "WHERE index=%s AND series=%s "
                      "AND tenor in ('3yr', '5yr', '7yr', '10yr')",
                      (index, series))
            start_date, = c.fetchone()
    else:
        start_date = None
    fh = logging.FileHandler(filename=os.path.join(os.getenv("LOG_DIR"), "index_curves.log"))
    formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
    fh.setFormatter(formatter)
    loggers = [logging.getLogger("analytics"), logging.getLogger("index_curves")]
    for logger in loggers:
        logger.setLevel(logging.INFO)
        logger.addHandler(fh)

    loggers[1].info(f"filling {index} {series}")
    g = calibrate_portfolio(index, series, ['3yr', '5yr', '7yr', '10yr'],
                            start_date)
    with conn.cursor() as c:
        for id, t in g:
            c.execute("INSERT INTO index_risk2 VALUES(%s, %s, %s, %s) ON CONFLICT (id) "
                      "DO UPDATE SET theta=%s, duration=%s, tweak=%s",
                      (id,) + tuple(t) + tuple(t))
    conn.commit()
    conn.close()