# packages
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from enum import Enum
import qis as qis
from qis import PerfParams, TimePeriod
import yfinance as yf


class LocalTests(Enum):
    LONG_IEF_SHORT_LQD = 1


def run_local_test(local_test: LocalTests):
    """Run local tests for development and debugging purposes.

    These are integration tests that download real data and generate reports.
    Use for quick verification during development.
    """

    perf_params = PerfParams(freq_drawdown='B', freq='B')
    kwargs = dict(fontsize=12, digits_to_show=1, sharpe_digits=2,
                  alpha_format='{0:+0.0%}',
                  beta_format='{:0.1f}',
                  perf_stats_labels=qis.PerfStatsLabels.TOTAL_DETAILED.value,
                  framealpha=0.9)

    time_period = TimePeriod('31Dec2021', '19Sep2022')

    if local_test == LocalTests.LONG_IEF_SHORT_LQD:
        prices = yf.download(tickers=['IEF', 'LQD', 'LQDH', 'IGIB'], start="2003-12-31", end=None, ignore_tz=True, auto_adjust=True)['Close']
        prices = time_period.locate(prices)
        rets = qis.to_returns(prices=prices, is_first_zero=True)
        navs1 = qis.returns_to_nav(returns=4.0*(rets.iloc[:, 0] - rets.iloc[:, 1]).rename('4x(Long 10y ETF / Short IG ETF)'))
        navs2 = qis.returns_to_nav(returns=-4.0*rets.iloc[:, 2].rename('4x(Short LQDH)'))
        navs3 = qis.returns_to_nav(returns=4.0*(0.5*rets.iloc[:, 0] - rets.iloc[:, 3]).rename('4x(Long 10y ETF / Short IGIB)'))
        spy = time_period.locate(yf.download(tickers=['SPY'], start="2003-12-31", end=None, ignore_tz=True, auto_adjust=True)['Close'])
        prices2 = pd.concat([navs1, navs2, navs3, spy], axis=1)

        with sns.axes_style('darkgrid'):
            fig1, axs = plt.subplots(2, 1, figsize=(15, 5), constrained_layout=True)
            qis.plot_prices_with_dd(prices=prices, perf_params=perf_params, axs=axs)

            fig2, ax = plt.subplots(1, 1, figsize=(12, 8), constrained_layout=True)
            qis.plot_prices(prices=prices2, perf_params=perf_params,
                            title='2022 YTD total performance of 4x(Long 10y ETF / Short IG ETF) vs SPY ETF',
                            ax=ax, **kwargs)

    plt.show()


if __name__ == '__main__':

    run_local_test(local_test=LocalTests.LONG_IEF_SHORT_LQD)
