from __future__ import annotations

from typing import Union, Optional, Literal

import pandas as pd
import numpy as np

import okama.common.helpers.helpers as helpers
import okama.portfolios.cashflow_strategies as cf
from okama import settings


def get_wealth_indexes_fv_with_cashflow(
    ror: Union[pd.Series, pd.DataFrame],
    portfolio_symbol: Optional[str],
    inflation_symbol: Optional[str],
    cashflow_parameters: cf.CashFlow,
    task: Literal["backtest", "monte_carlo"],
) -> Union[pd.Series, pd.DataFrame]:
    """
    Calculate wealth index Future Values (FV) for a series of returns with cash flows (withdrawals/contributions).

    Values of the wealth index correspond to the beginning of the month.
    """
    dcf_object = cashflow_parameters.parent.dcf
    dcf_object.cashflow_parameters = cashflow_parameters
    period_initial_amount = cashflow_parameters.initial_investment
    period_initial_amount_cached = period_initial_amount
    last_regular_cash_flow = 0
    amount = getattr(cashflow_parameters, "amount", None)
    if isinstance(ror, pd.DataFrame):
        portfolio_position = ror.columns.get_loc(portfolio_symbol)
    else:
        # for Series
        portfolio_position = 0
        ror = ror.to_frame()
    if cashflow_parameters.NAME == "CWD":
        drawdowns = helpers.Frame.get_drawdowns(ror.iloc[:, portfolio_position])
    cash_flow_ts = dcf_object.cashflow_parameters.time_series
    # check if iteration needed
    cashflow_iterate_condition = not (cash_flow_ts.empty or (cash_flow_ts == 0).all())
    if cashflow_iterate_condition:
        ror_cashflow_df = ror.assign(cashflow_ts=cash_flow_ts)
        ror_cashflow_df.fillna(0, inplace=True)
        n_rows = ror.shape[0]
        monthly_discount_rate = (1 + dcf_object.discount_rate) ** (1 / settings._MONTHS_PER_YEAR) - 1
        discount_factors = (1.0 + monthly_discount_rate) ** np.arange(n_rows)
        if task == "backtest":
            if dcf_object.cashflow_parameters.time_series_discounted_values:
                ror_cashflow_df.loc[:, "cashflow_ts"] = ror_cashflow_df.loc[:, "cashflow_ts"].mul(
                    discount_factors, axis=0
                )
        elif task == "monte_carlo":
            if not dcf_object.cashflow_parameters.time_series_discounted_values:
                ror_cashflow_df.loc[:, "cashflow_ts"] = ror_cashflow_df.loc[:, "cashflow_ts"].mul(
                    discount_factors, axis=0
                )
        else:
            raise ValueError(f"Unknown task: {task}. It must be 'monte_carlo' or 'backtest'")
    else:
        ror_cashflow_df = ror
        ror_cashflow_df.loc[:, "cashflow_ts"] = 0
    cash_flow_ts = ror_cashflow_df["cashflow_ts"]  # cash flow monthly time series
    periods_per_year = settings.frequency_periods_per_year[cashflow_parameters.frequency]
    if hasattr(cashflow_parameters, "indexation") and cashflow_parameters.frequency != "none":
        indexation_per_period = (1 + cashflow_parameters.indexation) ** (1 / periods_per_year) - 1
    if cashflow_parameters.frequency == "month" or cashflow_parameters.frequency == "none":
        # Fast Calculation
        s = pd.Series(dtype=float, name=portfolio_symbol)
        for n, row in enumerate(ror.itertuples()):
            date = row[0]
            r = row[portfolio_position + 1]
            if cashflow_parameters.frequency == "none":
                cashflow = 0
            elif cashflow_parameters.NAME == "fixed_amount":
                cashflow = amount * (1 + indexation_per_period) ** n
            elif cashflow_parameters.NAME == "fixed_percentage":
                cashflow = cashflow_parameters.percentage / periods_per_year * period_initial_amount
            elif cashflow_parameters.NAME == "time_series":
                cashflow = 0
            elif cashflow_parameters.NAME == "CWD":
                withdrawal_without_drawdowns = amount * (1 + indexation_per_period) ** n
                if drawdowns[date] < 0:
                    cashflow = cashflow_parameters._calculate_withdrawal_size(
                        drawdown=drawdowns[date],
                        withdrawal_without_drawdowns=withdrawal_without_drawdowns,
                    )
                else:
                    cashflow = withdrawal_without_drawdowns
            else:
                raise ValueError("Wrong cashflow strategy name value.")
            period_initial_amount = period_initial_amount * (r + 1) + cashflow + cash_flow_ts[date]
            date = row[0]
            s[date] = period_initial_amount
    elif cashflow_parameters.frequency != "month" and cashflow_parameters.frequency != "none":
        # Slow Calculation
        pandas_frequency = cashflow_parameters._pandas_frequency
        months_in_full_period = settings._MONTHS_PER_YEAR / cashflow_parameters.periods_per_year
        wealth_chunks = []  # Collect all chunks to concatenate once at the end
        for n, x in enumerate(ror_cashflow_df.resample(rule=pandas_frequency, convention="start")):
            ror_ts = x[1].iloc[:, portfolio_position]  # select ror part of the grouped data
            months_local = ror_ts.shape[0]
            period_fraction = months_local / months_in_full_period  # 1 for a full period
            cashflow_ts_local = x[1].loc[:, "cashflow_ts"].copy()
            last_date = ror_ts.index[-1]
            # CashFlow inside period (Extra withdrawals/contributions)
            if (cashflow_ts_local != 0).any():
                period_wealth_index = pd.Series(dtype=float, name=portfolio_symbol)
                for k, (date, r) in enumerate(ror_ts.items()):
                    if k == 0:
                        month_balance = period_initial_amount + cashflow_ts_local[date]
                    else:
                        month_balance = month_balance * (r + 1) + cashflow_ts_local[date]
                    period_wealth_index[date] = month_balance
            else:
                period_wealth_index = period_initial_amount * (1 + ror_ts).cumprod()
            # CashFlow END period
            if cashflow_parameters.NAME == "fixed_amount":
                cashflow_value = amount * (1 + indexation_per_period) ** n
            elif cashflow_parameters.NAME == "fixed_percentage":
                cashflow_value = cashflow_parameters.percentage / periods_per_year * period_initial_amount
            elif cashflow_parameters.NAME == "VDS":
                cashflow_value = cashflow_parameters._calculate_withdrawal_size(
                    last_withdrawal=last_regular_cash_flow if n > 0 else 0,
                    balance=period_initial_amount,
                    number_of_periods=n,
                )
            elif cashflow_parameters.NAME == "CWD":
                withdrawal_without_drawdowns = amount * (1 + indexation_per_period) ** n
                if drawdowns[last_date] < 0:
                    cashflow_value = cashflow_parameters._calculate_withdrawal_size(
                        drawdown=drawdowns[last_date],
                        withdrawal_without_drawdowns=withdrawal_without_drawdowns,
                    )
                else:
                    cashflow_value = withdrawal_without_drawdowns
            else:
                raise ValueError("Wrong cashflow_method value.")
            cashflow_value *= period_fraction  # adjust cash flow to the period length (months)
            period_final_balance = period_wealth_index.iloc[-1] + cashflow_value
            period_wealth_index.iloc[-1] = period_final_balance
            period_initial_amount = period_final_balance
            wealth_chunks.append(period_wealth_index)
        wealth_df = (
            pd.concat(wealth_chunks, sort=False)
            if wealth_chunks
            else pd.DataFrame(dtype=float, columns=[portfolio_symbol])
        )
        s = wealth_df.squeeze()
    first_date = s.index[0]
    first_wealth_index_date = first_date - 1  # set first date to one month earlie
    s.loc[first_wealth_index_date] = period_initial_amount_cached
    if inflation_symbol:
        cum_inflation = helpers.Frame.get_wealth_indexes(
            ror=ror.loc[:, inflation_symbol], initial_amount=period_initial_amount_cached
        )
        wealth_index = pd.concat([s, cum_inflation], axis="columns")
    else:
        wealth_index = s
    wealth_index.sort_index(ascending=True, inplace=True)
    return wealth_index


def get_cash_flow_fv(
    ror: Union[pd.Series, pd.DataFrame],
    portfolio_symbol: Optional[str],
    cashflow_parameters: cf.CashFlow,
    task: Literal["backtest", "monte_carlo"],
) -> Union[pd.Series, pd.DataFrame]:
    """
    Calculate cash flow future values (FV) for a series of returns according to withdrawal/contributions strategies.
    """
    dcf_object = cashflow_parameters.parent.dcf
    dcf_object.cashflow_parameters = cashflow_parameters
    period_initial_amount = cashflow_parameters.initial_investment
    last_regular_cash_flow = 0
    cs_fv = pd.Series(dtype=float, name="cash_flow_fv")
    amount = getattr(cashflow_parameters, "amount", None)
    if cashflow_parameters.NAME == "CWD":
        drawdowns = helpers.Frame.get_drawdowns(ror)
    if isinstance(ror, pd.DataFrame):
        portfolio_position = ror.columns.get_loc(portfolio_symbol)
    else:
        # for Series
        portfolio_position = 0
        ror = ror.to_frame()
    cash_flow_ts = dcf_object.cashflow_parameters.time_series
    # check if iteration needed
    cashflow_iterate_condition = not (cash_flow_ts.empty or (cash_flow_ts == 0).all())
    if cashflow_iterate_condition:
        ror_cashflow_df = ror.assign(cashflow_ts=cash_flow_ts)
        ror_cashflow_df.fillna(0, inplace=True)
        n_rows = ror.shape[0]
        monthly_discount_rate = (1 + dcf_object.discount_rate) ** (1 / settings._MONTHS_PER_YEAR) - 1
        discount_factors = (1.0 + monthly_discount_rate) ** np.arange(n_rows)
        if task == "backtest":
            if dcf_object.cashflow_parameters.time_series_discounted_values:
                ror_cashflow_df.loc[:, "cashflow_ts"] = ror_cashflow_df.loc[:, "cashflow_ts"].mul(
                    discount_factors, axis=0
                )
        elif task == "monte_carlo":
            if not dcf_object.cashflow_parameters.time_series_discounted_values:
                ror_cashflow_df.loc[:, "cashflow_ts"] = ror_cashflow_df.loc[:, "cashflow_ts"].mul(
                    discount_factors, axis=0
                )
        else:
            raise ValueError(f"Unknown task: {task}. It must be 'monte_carlo' or 'backtest'")
    else:
        ror_cashflow_df = ror
        ror_cashflow_df.loc[:, "cashflow_ts"] = 0
    cash_flow_ts = ror_cashflow_df["cashflow_ts"]  # cash flow monthly time series
    periods_per_year = settings.frequency_periods_per_year[cashflow_parameters.frequency]
    if hasattr(cashflow_parameters, "indexation") and cashflow_parameters.frequency != "none":
        indexation_per_period = (1 + cashflow_parameters.indexation) ** (1 / periods_per_year) - 1
    if cashflow_parameters.frequency == "month" or cashflow_parameters.frequency == "none":
        # Fast Calculation
        for n, row in enumerate(ror.itertuples()):
            date = row[0]
            r = row[portfolio_position + 1]
            # Calculate regular cash flow
            if cashflow_parameters.frequency == "none":
                cashflow = 0
            elif cashflow_parameters.NAME == "fixed_amount":
                cashflow = amount * (1 + indexation_per_period) ** n
            elif cashflow_parameters.NAME == "fixed_percentage":
                cashflow = cashflow_parameters.percentage / periods_per_year * period_initial_amount
            elif cashflow_parameters.NAME == "time_series":
                cashflow = 0
            elif cashflow_parameters.NAME == "CWD":
                withdrawal_without_drawdowns = amount * (1 + indexation_per_period) ** n
                if drawdowns[date] < 0:
                    cashflow = cashflow_parameters._calculate_withdrawal_size(
                        drawdown=drawdowns[date],
                        withdrawal_without_drawdowns=withdrawal_without_drawdowns,
                    )
                else:
                    cashflow = withdrawal_without_drawdowns
            else:
                raise ValueError("Wrong cashflow strategy name value.")
            # add Extra Withdrawals/Contributions
            cs_value = cashflow + cash_flow_ts[date]
            period_initial_amount = period_initial_amount * (r + 1) + cs_value
            cs_fv[date] = cs_value
    elif cashflow_parameters.frequency != "month" and cashflow_parameters.frequency != "none":
        # Slow Calculation
        pandas_frequency = settings.frequency_mapping[cashflow_parameters.frequency]
        months_in_full_period = settings._MONTHS_PER_YEAR / cashflow_parameters.periods_per_year
        cashflow_chunks = []  # Collect all chunks to concatenate once at the end
        for n, x in enumerate(ror_cashflow_df.resample(rule=pandas_frequency, convention="start")):
            ror_ts = x[1].iloc[:, portfolio_position]  # select ror part of the grouped data
            months_local = ror_ts.shape[0]
            period_fraction = months_local / months_in_full_period  # 1 for a full period
            cashflow_ts_local = x[1].loc[:, "cashflow_ts"].copy()
            last_date = ror_ts.index[-1]
            # CashFlow inside period (Extra cash flow)
            if (cashflow_ts_local != 0).any():
                period_wealth_index = pd.Series(dtype=float, name=portfolio_symbol)
                for k, (date, r) in enumerate(ror_ts.items()):
                    if k == 0:
                        month_balance = period_initial_amount
                    else:
                        month_balance = month_balance * (r + 1) + cashflow_ts_local[date]
                    period_wealth_index[date] = month_balance
            else:
                period_wealth_index = period_initial_amount * (1 + ror_ts).cumprod()
            # CashFlow END period (Regular Cash Flow)
            if cashflow_parameters.NAME == "fixed_amount":
                cashflow_value = amount * (1 + indexation_per_period) ** n
            elif cashflow_parameters.NAME == "fixed_percentage":
                cashflow_value = cashflow_parameters.percentage / periods_per_year * period_initial_amount
            elif cashflow_parameters.NAME == "VDS":
                cashflow_value = cashflow_parameters._calculate_withdrawal_size(
                    last_withdrawal=last_regular_cash_flow if n > 0 else 0,
                    balance=period_initial_amount,
                    number_of_periods=n,
                )
            elif cashflow_parameters.NAME == "CWD":
                withdrawal_without_drawdowns = amount * (1 + indexation_per_period) ** n
                if drawdowns[last_date] < 0:
                    cashflow_value = cashflow_parameters._calculate_withdrawal_size(
                        drawdown=drawdowns[last_date],
                        withdrawal_without_drawdowns=withdrawal_without_drawdowns,
                    )
                else:
                    cashflow_value = withdrawal_without_drawdowns
            else:
                raise ValueError("Wrong cashflow_method value.")
            cashflow_value *= period_fraction  # adjust cash flow to the period length (months)
            last_regular_cash_flow = cashflow_value
            period_final_balance = period_wealth_index.iloc[-1] + cashflow_value
            period_wealth_index.iloc[-1] = period_final_balance
            period_initial_amount = period_final_balance
            cashflow_ts_local.iloc[-1] += cashflow_value
            cashflow_chunks.append(cashflow_ts_local)
        cs_fv = pd.concat(cashflow_chunks, sort=False) if cashflow_chunks else cs_fv
    return cs_fv


def remove_negative_values(input_s: pd.Series) -> pd.Series:
    if not isinstance(input_s, pd.Series):
        raise TypeError("input_s must be a pd.Series")
    s = input_s.copy()
    condition = s <= 0
    try:
        survival_date = s[condition].index[0]
        s[survival_date] = 0
        s[s.index > survival_date] = np.nan
    except IndexError:
        pass
    return s


def discount_monthly_cash_flow(
    cash_flow_fv: Union[pd.Series, pd.DataFrame],
    annual_effective_discount_rate: float,
    reverse: bool = False,
) -> Union[pd.Series, pd.DataFrame]:
    number_of_months = cash_flow_fv.shape[0]
    monlthly_discount_rate = (1 + annual_effective_discount_rate) ** (1 / settings._MONTHS_PER_YEAR) - 1
    if not reverse:
        discount_factors = (1.0 + monlthly_discount_rate) ** np.arange(number_of_months)
    else:
        discount_factors = (1.0 + monlthly_discount_rate) ** np.arange(number_of_months)[::-1]
    cash_flow_pv = cash_flow_fv.div(discount_factors, axis=0)
    return cash_flow_pv
