import pandas as pd
from serenitas.analytics.api import CreditIndex, Swaption, BlackSwaption
import datetime

from serenitas.utils.db2 import dbconn
from contextlib import contextmanager
from itertools import starmap
from functools import partial
from multiprocessing import Pool


def get_data(conn, index, series, date=datetime.date.min):
    df = pd.read_sql_query(
        "SELECT * from swaption_ref_quotes JOIN swaption_quotes "
        "USING (ref_id) WHERE index=%s and series=%s "
        "and quotedate >=%s ORDER BY quotedate",
        conn,
        params=(index, series, date),
        parse_dates=["quotedate", "expiry"],
    )
    df.loc[
        (df.quote_source == "GS") & (df["index"] == "HY"),
        ["pay_bid", "pay_offer", "rec_bid", "rec_offer"],
    ] *= 100
    if not df.empty:
        df.quotedate = df.quotedate.dt.tz_convert("America/New_York")
    return df


def get_data_latest(conn):
    df = pd.read_sql_query(
        "SELECT quotedate, index, series, expiry, ref, "
        "quote_source, swaption_quotes.* "
        "FROM swaption_ref_quotes "
        "JOIN swaption_quotes USING (ref_id) "
        "LEFT JOIN swaption_calib USING (quote_id) "
        "WHERE swaption_calib.quote_id is NULL",
        conn,
        parse_dates=["quotedate", "expiry"],
    )
    df.loc[
        (df.quote_source == "GS") & (df["index"] == "HY"),
        ["pay_bid", "pay_offer", "rec_bid", "rec_offer"],
    ] *= 100
    if not df.empty:
        df.quotedate = df.quotedate.dt.tz_convert("America/New_York")
    conn.commit()
    return df


def calib(option, ref, strike, pay_bid, pay_offer, rec_bid, rec_offer):
    option.ref = ref
    option.strike = strike
    r = []
    logger.debug(
        f"{ref=}\t{strike=}\t{pay_bid=}\t{pay_offer=}\t{rec_bid=}\t{rec_offer=}"
    )
    for price_type in ["price_black", "price"]:
        for option_type in ["pay", "rec"]:
            if option_type == "pay":
                mid = (pay_bid + pay_offer) / 2 * 1e-2
                option.option_type = "payer"
            else:
                mid = (rec_bid + rec_offer) / 2 * 1e-2
                option.option_type = "receiver"
            if mid == 0.0:
                logger.info("0. mid, skipping.")
                r.append(0.0)
                continue
            try:
                if price_type == "price_black":
                    BlackSwaption.price.fset(option, mid)
                else:
                    setattr(option, price_type, mid)
            except (ValueError, SystemError) as e:
                if "Failed" in str(e):
                    logger.error(e)
                    logger.error("probably data error")
                    logger.info(
                        f"ref: {ref}, strike: {strike}, expiry: {option.exercise_date}, mid: {mid}"
                    )
                r.append(None)
            else:
                r.append(option.sigma)
    return r


@contextmanager
def MaybePool(nproc):
    yield Pool(nproc) if nproc > 1 else None


def calibrate(index_type=None, series=None, date=None, nproc=4, latest=False):
    sql_str = "INSERT INTO swaption_calib VALUES({}) ON CONFLICT DO NOTHING".format(
        ",".join(["%s"] * 5)
    )
    serenitasdb = dbconn(
        "serenitasdb", extra_option="-c idle_in_transaction_session_timeout=10s"
    )
    if latest:
        data = get_data_latest(serenitasdb)
    else:
        data = get_data(serenitasdb, index_type, series, date)

    with MaybePool(nproc) as pool:
        pstarmap = pool.starmap if pool else starmap
        for k, v in data.groupby([data["quotedate"].dt.date, "index", "series"]):
            trade_date, index_type, series = k
            logger.debug(f"{trade_date} {index_type}{series}")
            index = CreditIndex(index_type, series, "5yr", trade_date)
            for expiry, df in v.groupby(["expiry"]):
                logger.debug(expiry)
                try:
                    option = Swaption(index, expiry.date(), 100)
                except ValueError as e:
                    logger.error(e)
                    continue
                mycalib = partial(calib, option)
                r = pstarmap(
                    mycalib,
                    df[
                        [
                            "ref",
                            "strike",
                            "pay_bid",
                            "pay_offer",
                            "rec_bid",
                            "rec_offer",
                        ]
                    ]
                    .sort_values("strike")
                    .itertuples(index=False, name=None),
                )
                to_insert = [[a] + b for a, b in zip(df.quote_id, r)]
                with serenitasdb.cursor() as c:
                    c.executemany(sql_str, to_insert)
                serenitasdb.commit()


if __name__ == "__main__":
    import logging
    from serenitas.utils import SerenitasFileHandler

    logger = logging.getLogger("swaption_calib")
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--index", required=False, type=lambda s: s.upper(), dest="index_type"
    )
    parser.add_argument("--series", required=False, type=int, default=28)
    parser.add_argument("--date", required=False, default=datetime.date.min)
    parser.add_argument("--latest", required=False, action="store_true")
    parser.add_argument("--nproc", required=False, type=int, default=4)
    parser.add_argument(
        "-d", "--debug", action="store_true", help="more verbose logging"
    )
    args = parser.parse_args()

    logger.setLevel(logging.DEBUG if args.debug else logging.INFO)
    if not args.debug:
        handler = SerenitasFileHandler(f"calib_swaptions_{datetime.date.today()}.log")
    else:
        handler = logging.StreamHandler()
        handler.setFormatter(SerenitasFileHandler._formatter)
    if not logger.handlers:
        logger.addHandler(handler)
    if args.latest:
        calibrate(latest=True, nproc=args.nproc)
    else:
        calibrate(**vars(args))