import pandas as pd
from pandas.tseries.offsets import BDay

from serenitas.utils.db import dbconn, serenitas_engine

dawndb = dbconn("dawndb")
df_balances = pd.read_sql_query(
    "SELECT * FROM strategy_im WHERE fund='SERCGMAST'",
    dawndb,
    parse_dates=["date"],
    index_col=["date"],
).sort_index()
df_balances[["broker", "strategy"]] = df_balances[["broker", "strategy"]].astype(
    "category"
)


def get_rates(broker):
    rate_index = "SOFRRATE" if broker == "BARCLAYS" else "FED_FUND"

    return pd.read_sql_query(
        "SELECT date, rate FROM rates where name=%s",
        serenitas_engine,
        params=(rate_index,),
        parse_dates=["date"],
        index_col=["date"],
    ).sort_index()


def f(broker, start_date, end_date):
    df_rates = get_rates(broker)
    df = (
        df_balances[df_balances.broker == broker]
        .set_index("strategy", append=True)["amount"]
        .unstack("strategy")
    )
    df[df.isnull()] = 0.0
    drange = pd.date_range(pd.Timestamp(start_date) - BDay(), end_date)
    rates = df_rates.reindex(drange, method="ffill") / 100 / 360
    df = df.reindex(drange, method="ffill")
    if broker in ["BAML_ISDA", "CITI"]:
        d = {}
        for strat in df:
            s = df.loc[start_date:, strat]
            ir_bal = 0.0
            for bal, r in zip(s.values, rates.loc[start_date:, "rate"].values):
                bal += ir_bal
                ir_bal += bal * r
            d[strat] = ir_bal
        return pd.Series(d, name="amount").to_frame()
    else:
        return (
            (df.loc[start_date:] * rates.loc[start_date:].values)
            .sum()
            .to_frame(name="amount")
        )


def export_data(start, end):
    dfs = {}
    for cp in ("GS", "MS", "BAML_ISDA", "CITI", "CS", "BNP", "JPM", "BARCLAYS"):
        dfs[cp] = f(cp, start, end)
    df = pd.concat(dfs, names=["broker", "folder"])
    df = df[df.amount != 0.0]
    df.amount *= -1.0
    return df


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("start_date")
    parser.add_argument("end_date")
    args = parser.parse_args()
    df = export_data(args.start_date, args.end_date)