"""
create total return blended portfolios based on weights
"""

# packages
import numpy as np
import pandas as pd
from numba import njit
from enum import Enum
from typing import Union, Dict, Tuple, List, Optional
# qis
import qis.utils as qu
from qis.portfolio.portfolio_data import PortfolioData


def backtest_model_portfolio(prices: pd.DataFrame,
                             weights: Union[Dict[str, float], List[float], np.ndarray, pd.DataFrame, pd.Series],
                             rebalancing_freq: Optional[str] = 'QE',
                             initial_nav: float = 100,
                             funding_rate: pd.Series = None,  # annualised on positive / negative cash balances
                             management_fee: float = None,  # annualised
                             instruments_carry: pd.DataFrame = None,  # on nav
                             rebalancing_costs: Union[float, pd.Series] = None,  # annualised rebalancing costs in bp
                             weight_implementation_lag: Optional[int] = None,  # applies for weight is pd.Dataframe
                             constant_trade_level: float = None,
                             is_rebalanced_at_first_date: bool = False,
                             ticker: str = 'Portfolio'  # default ticker
                             ) -> PortfolioData:
    """
    simulate portfolio given prices and weights
    include_start_date if index rebalanced at start date
    the safest weight is to pass weights as Dict or pd.Dataframe - this enforces the alignment with prices
    does not rebalancing_freq when dates are pd.DataFrame
    funding_rate is funding rate on cash annualised
    management_fee is man fee on strategy nav annualised
    """
    if not isinstance(prices, pd.DataFrame):
        raise ValueError(f"prices type={type(prices)} must be pd.Dataframe")
    if isinstance(weights, pd.Series):  # map to dict
        weights = weights.to_dict()

    if isinstance(weights, Dict):  # map to np
        qu.assert_list_subset(large_list=prices.columns.to_list(),
                              list_sample=list(weights.keys()),
                              message=f"weights columns must be aligned with price columns")
        weights = prices.columns.map(weights).to_numpy()
    elif isinstance(weights, List):
        weights = np.array(weights)

    # align weights with prices
    if isinstance(weights, np.ndarray):
        if weights.shape[0] != len(prices.columns):
            raise ValueError(f"number of weights must be aligned with number of price columns")
        if len(weights.shape) > 1:
            raise ValueError(f"only single aray is allowed")

        is_rebalancing = qu.generate_rebalancing_indicators(df=prices,
                                                            freq=rebalancing_freq,
                                                            include_start_date=is_rebalanced_at_first_date)

        portfolio_rebalance_dates = is_rebalancing[is_rebalancing == True]
        portfolio_weights = pd.DataFrame(data=qu.repeat_by_rows(weights, n=len(portfolio_rebalance_dates)),
                                         index=portfolio_rebalance_dates,
                                         columns=prices.columns)

    elif isinstance(weights, pd.DataFrame):
        qu.assert_list_subset(large_list=prices.columns.to_list(),
                              list_sample=weights.columns.to_list(),
                              message=f"weights columns must be aligned with price columns")
        if prices.index[0] > weights.index[0]:
            raise ValueError(f"price dates {prices.index[0]} are after weights start date {weights.index[0]}")
        portfolio_weights = weights[prices.columns]  # alighn

        # implementation lag is only valid for quant-generated weights
        if weight_implementation_lag is not None and weight_implementation_lag > 0:
            rebalancing_index = portfolio_weights.index + pd.Timedelta(days=weight_implementation_lag)
        else:
            rebalancing_index = portfolio_weights.index
        is_rebalancing = qu.set_rebalancing_timeindex_on_given_timeindex(given_index=prices.index,
                                                                         rebalancing_index=rebalancing_index)
    else:
        raise NotImplementedError(f"unsupported weights type = {type(weights)}")

    # adjust rates at rebealncing
    if funding_rate is not None:
        funding_rate_dt = qu.multiply_df_by_dt(df=funding_rate, dates=prices.index, lag=0)
    else:
        funding_rate_dt = pd.Series(0.0, index=prices.index)

    if management_fee is not None:
        management_fee_dt = qu.multiply_df_by_dt(df=pd.Series(management_fee, index=prices.index), dates=prices.index, lag=0)
    else:
        management_fee_dt = pd.Series(0.0, index=prices.index)

    if instruments_carry is not None:
        instruments_carry_dt = qu.multiply_df_by_dt(df=instruments_carry, dates=prices.index, lag=0)
    else:
        instruments_carry_dt = pd.Series(0.0, index=prices.index)

    if rebalancing_costs is not None:
        if isinstance(rebalancing_costs, pd.Series):
            rebalancing_costs = rebalancing_costs[prices.columns].to_numpy()

    nav, units, effective_weights, realized_costs = backtest_rebalanced_portfolio(prices=prices.to_numpy(),
                                                                                  weights=portfolio_weights.to_numpy(),
                                                                                  is_rebalancing=is_rebalancing.to_numpy(),
                                                                                  funding_rate_dt=funding_rate_dt.to_numpy(),
                                                                                  management_fee_dt=management_fee_dt.to_numpy(),
                                                                                  instruments_carry_dt=instruments_carry_dt.to_numpy(),
                                                                                  initial_nav=initial_nav,
                                                                                  constant_trade_level=constant_trade_level,
                                                                                  rebalancing_costs=rebalancing_costs)

    portfolio_nav = pd.Series(nav, index=prices.index)
    if ticker is not None:
        portfolio_nav = portfolio_nav.rename(ticker)

    output_portfolio_data = PortfolioData(nav=portfolio_nav,
                                          units=pd.DataFrame(units, index=prices.index, columns=prices.columns),
                                          weights=pd.DataFrame(effective_weights, index=prices.index, columns=prices.columns),
                                          input_weights=weights,
                                          is_rebalancing=is_rebalancing,
                                          prices=prices,
                                          realized_costs=pd.DataFrame(realized_costs, index=prices.index, columns=prices.columns),
                                          ticker=ticker)

    return output_portfolio_data


@njit
def backtest_rebalanced_portfolio(prices: np.ndarray,
                                  weights: np.ndarray,
                                  is_rebalancing: np.ndarray,
                                  funding_rate_dt: np.ndarray = None,
                                  management_fee_dt: np.ndarray = None,
                                  instruments_carry_dt: np.ndarray = None,
                                  initial_nav: float = 100.0,
                                  constant_trade_level: float = None,
                                  rebalancing_costs: Union[float, np.ndarray] = None  # proportional rebalancing costs
                                  ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    if prices.shape[0] != is_rebalancing.shape[0]:
        raise ValueError(f"prices.shape[0] != is_rebalancing.shape[0]")
    if weights.shape[1] != prices.shape[1]:
        raise ValueError(f"weights.shape[1] != prices.shape[1]")

    if funding_rate_dt is None:
        funding_rate_dt = np.zeros(prices.shape[0])

    if management_fee_dt is None:
        management_fee_dt = np.zeros(prices.shape[0])

    # initialize
    current_rebalancing_idx = 0
    units = np.zeros_like(prices)
    realized_costs = np.zeros_like(prices)
    nav = np.zeros(prices.shape[0])
    cash_balances = np.zeros(prices.shape[0])

    # build initial portfolio
    current_prices = prices[0, :]
    if is_rebalancing[0]:
        if constant_trade_level is not None:
            current_units = (constant_trade_level * weights[current_rebalancing_idx]) / current_prices
        else:
            current_units = (initial_nav * weights[current_rebalancing_idx]) / current_prices

        current_units[np.isnan(current_units)] = 0
        current_cash_balance = initial_nav - np.nansum(current_units * current_prices)
        current_rebalancing_idx += 1
    else:
        current_units = np.zeros(prices.shape[1])
        current_cash_balance = initial_nav
    units[0, :] = current_units
    nav[0] = np.nansum(current_units * current_prices) + current_cash_balance  # need to be adjusted when cost are present for is_rebalancing[0] = True
    cash_balances[0] = current_cash_balance

    # loop over t
    for t in np.arange(1, prices.shape[0]):
        current_units = units[t - 1]
        current_prices = prices[t, :]
        management_fee = management_fee_dt[t]*nav[t-1]
        current_cash_balance = cash_balances[t-1] * (1.0 + funding_rate_dt[t]) - management_fee
        
        if instruments_carry_dt is not None:
            carry = np.nansum(current_units * current_prices * instruments_carry_dt[t])
            current_cash_balance += carry

        if is_rebalancing[t]:
            if constant_trade_level:
                current_nav0 = constant_trade_level
            else:
                current_nav0 = np.nansum(current_units * current_prices) + current_cash_balance
            current_units = (current_nav0 * weights[current_rebalancing_idx]) / current_prices
            current_units[np.isnan(current_units)] = 0
            units_change = current_units - units[t-1]
            change_in_cash = -np.nansum(units_change*current_prices)
            if rebalancing_costs is not None:
                realized_costs_t = rebalancing_costs*current_prices*np.abs(units_change)
                realized_costs[t, :] = realized_costs_t
                change_in_cash -= np.nansum(realized_costs_t)
            current_cash_balance = current_cash_balance + change_in_cash
            current_rebalancing_idx += 1

        # store
        units[t, :] = current_units
        nav[t] = np.nansum(current_units * current_prices) + current_cash_balance
        cash_balances[t] = current_cash_balance

    effective_weights = np.divide(units * prices, qu.repeat_by_columns(a=nav, n=prices.shape[1]))

    return nav, units, effective_weights, realized_costs


class LocalTests(Enum):
    BLENDED = 1
    COSTS = 2


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.
    """

    import matplotlib.pyplot as plt
    import qis.plots.derived.prices as ppd

    from qis.test_data import load_etf_data
    prices = load_etf_data().dropna()

    prices = prices[['SPY', 'TLT']]
    # prices.iloc[:200, :] = np.nan
    print(prices)
    
    if local_test == LocalTests.BLENDED:

        portfolio_nav_1_0 = backtest_model_portfolio(prices=prices,
                                                     weights=np.array([1.0, 0.0]),
                                                     rebalancing_freq='QE').get_portfolio_nav()

        portfolio_nav_5_5 = backtest_model_portfolio(prices=prices,
                                                     weights=np.array([1.0, 0.5]),
                                                     rebalancing_freq='QE').get_portfolio_nav()

        portfolio_nav_0_1 = backtest_model_portfolio(prices=prices,
                                                     weights=np.array([1.0, 1.0]),
                                                     rebalancing_freq='QE').get_portfolio_nav()

        portfolio_nav = pd.concat([portfolio_nav_1_0, portfolio_nav_5_5, portfolio_nav_0_1], axis=1)
        portfolio_nav.columns = ['x1=100, x2=0', 'x1=100, x2=50', 'x1=100, x2=100']
        print(portfolio_nav)
        fig, ax = plt.subplots(1, 1, figsize=(12, 10))
        ppd.plot_prices(prices=portfolio_nav, ax=ax)

    elif local_test == LocalTests.COSTS:
        portfolio_nav = backtest_model_portfolio(prices=prices,
                                                 weights=np.array([1.0, 1.0]),
                                                 rebalancing_freq='QE')

        portfolio_nav.plot_pnl()

    plt.show()


if __name__ == '__main__':

    run_local_test(local_test=LocalTests.BLENDED)
