from typing import Dict, Optional, List, Any, Type, Union
from abc import ABC, abstractmethod
from joblib import Parallel, delayed

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.axes import Axes

from okama import macro, asset, settings
from okama.common import validators
from okama.common.helpers import helpers
from okama.common.error import ShortPeriodLengthError


class ListMaker(ABC):
    """
    Abstract base class for building an asset list.

    Parameters
    ----------
    assets : list, default None
        List of assets. Could include tickers or asset-like objects (`Asset`, `Portfolio`).
        If None, a single-asset list with a default ticker is used.
    first_date : str, default None
        First date of monthly return time series.
        If None, the first date is calculated automatically as the oldest available date for the listed assets.
    last_date : str, default None
        Last date of monthly return time series.
        If None, the last date is calculated automatically as the newest available date for the listed assets.
    ccy : str, default 'USD'
        Base currency for the list of assets. All risk metrics and returns are adjusted to the base currency.
    inflation : bool, default True
        Defines whether to take inflation data into account in the calculations.
        Including inflation could limit available data (`first_date`, `last_date`)
        as the inflation data is usually published with a one-month delay.
        With `inflation=False`, some properties like real return are not available.
    """

    def __init__(
        self,
        assets: Optional[List[Union[str, Type]]] = None,
        *,
        first_date: Optional[str] = None,
        last_date: Optional[str] = None,
        ccy: str = "USD",
        inflation: bool = True,
    ):
        self._assets = list(dict.fromkeys(assets)) if assets else [settings.default_ticker]
        self._currency = asset.Asset(symbol=f"{ccy}.FX", first_date=first_date, last_date=last_date)
        self.asset_obj_dict = ListMaker._get_asset_obj_dict(
            self._list_of_asset_like_objects,
            first_date=first_date,
            last_date=last_date,
        )
        (
            self.first_date,
            self.last_date,
            self.newest_asset,
            self.eldest_asset,
            self.names,
            self.currencies,
            self.assets_first_dates,
            self.assets_last_dates,
            self._assets_ror,
        ) = self._make_list(first_date=first_date, last_date=last_date).values()
        if first_date:
            self.first_date = max(self.first_date, pd.to_datetime(first_date))
        self._assets_ror = self._assets_ror[self.first_date :]
        if last_date:
            self.last_date = min(self.last_date, pd.to_datetime(last_date))
        if inflation:
            self.inflation: str = f"{ccy}.INFL"
            self._inflation_instance = macro.Inflation(self.inflation, first_date, last_date)
            self.inflation_first_date: pd.Timestamp = self._inflation_instance.first_date
            self.inflation_last_date: pd.Timestamp = self._inflation_instance.last_date
            self.first_date = max(self.first_date, self.inflation_first_date)
            self.last_date = min(self.last_date, self.inflation_last_date)
            self.inflation_ts: pd.Series = self._inflation_instance.values_monthly.loc[self.first_date : self.last_date]
            # Add inflation to the date range dict
            self.assets_first_dates.update({self.inflation: self.inflation_first_date})
            self.assets_last_dates.update({self.inflation: self.inflation_last_date})
        self._assets_ror: pd.DataFrame = self._assets_ror[self.first_date : self.last_date]
        self.period_length: float = round((self.last_date - self.first_date) / np.timedelta64(365, "D"), ndigits=2)
        self.pl = settings.PeriodLength(
            self._assets_ror.shape[0] // settings._MONTHS_PER_YEAR,
            self._assets_ror.shape[0] % settings._MONTHS_PER_YEAR,
        )
        self._pl_txt = f"{self.pl.years} years, {self.pl.months} months"
        self._dividend_yield: pd.DataFrame = pd.DataFrame(dtype=float)
        self._assets_dividends_ts: pd.DataFrame = pd.DataFrame(dtype=float)

    def __iter__(self):
        return iter(self.asset_obj_dict.values())

    @abstractmethod
    def __repr__(self):
        pass

    def __len__(self):
        """
        Return the number of assets in the list.

        Returns
        -------
        int
            Number of assets.
        """
        return len(self.symbols)

    def __getitem__(self, item):
        """
        Return an asset object by index.

        Parameters
        ----------
        item : int
            Index of the asset.

        Returns
        -------
        Asset
            Asset object.
        """
        return list(self.asset_obj_dict.values())[item]

    @staticmethod
    def _get_asset_obj_dict(
        ls: list,
        first_date: Optional[str] = None,
        last_date: Optional[str] = None,
    ) -> dict:
        """
        Create a dictionary of Asset objects from a list of symbols or Asset-like objects.

        Parameters
        ----------
        ls : list
            List of symbols or Asset-like objects.
        first_date : str, optional
            First date for the assets.
        last_date : str, optional
            Last date for the assets.

        Returns
        -------
        dict
            Dictionary of Asset objects with symbols as keys.

        Raises
        ------
        ShortPeriodLengthError
            If an asset's historical data period is too short (less than 3 months).
        """

        def get_item(symbol):
            asset_item = (
                symbol
                if hasattr(symbol, "symbol") and hasattr(symbol, "ror")
                else asset.Asset(symbol, first_date=first_date, last_date=last_date)
            )
            # Primary guard: use asset_item.pl when available
            try:
                if asset_item.pl.years == 0 and asset_item.pl.months <= 2:
                    raise ShortPeriodLengthError(
                        f"{asset_item.symbol} period length is {asset_item.pl.months}. It should be at least 3 months."
                    )
            except AttributeError:
                # Fallback guard: if pl is missing or malformed, use the length of ror series
                pass

            # Additional robust guard: ensure at least 3 monthly observations in ror
            try:
                ror_len = len(asset_item.ror)
            except Exception:
                ror_len = 0
            # Only enforce if pl did not already trigger; allow both to be consistent
            if ror_len < 3:
                # Try to provide a meaningful months value in the message
                months_msg = getattr(getattr(asset_item, "pl", None), "months", ror_len)
                raise ShortPeriodLengthError(
                    f"{asset_item.symbol} period length is {months_msg}. It should be at least 3 months."
                )
            return asset_item

        # For a single item avoid parallel overhead and ensure exceptions propagate directly
        if len(ls) == 1:
            obj = get_item(ls[0])
            return {obj.symbol: obj}
        asset_obj_list = Parallel(n_jobs=-1, backend="threading")(delayed(get_item)(s) for s in ls)
        return {obj.symbol: obj for obj in asset_obj_list}

    def _make_list(self, first_date: Optional[str], last_date: Optional[str]) -> dict:
        """
        Make an asset list from a list of symbols.

        Parameters
        ----------
        first_date : str, optional
            First date of monthly return time series.
        last_date : str, optional
            Last date of monthly return time series.

        Returns
        -------
        dict
            Dictionary containing various properties of the asset list,
            such as first_date, last_date, asset names, currencies, etc.
        """
        base_currency_ticker: str = self._currency.ticker
        currency_first_date: pd.Timestamp = self._currency.first_date
        currency_last_date: pd.Timestamp = self._currency.last_date

        own_first_dates: Dict[str, pd.Timestamp] = {}
        own_last_dates: Dict[str, pd.Timestamp] = {}
        first_dates: Dict[str, pd.Timestamp] = {}
        last_dates: Dict[str, pd.Timestamp] = {}
        names: Dict[str, str] = {}
        currencies: Dict[str, str] = {}
        input_first_date = pd.to_datetime(first_date) if first_date else None
        input_last_date = pd.to_datetime(last_date) if last_date else None

        # Collect all rate of return series first, then concatenate once (more efficient)
        ror_series_list: List[pd.Series] = []
        for asset_item in self.asset_obj_dict.values():
            # get asset own first and last dates
            asset_own_first_date = asset_item.first_date
            asset_own_last_date = asset_item.last_date

            ror_series = self._make_ror(asset_item, base_currency_ticker)
            ror_series_list.append(ror_series)

            # get asset first and last dates after adjusting to the currency
            asset_first_date = ror_series.index[0].to_timestamp(how="start")
            asset_last_date = ror_series.index[-1].to_timestamp(how="start")

            # check first and last dates
            fd_max = max(x for x in [asset_first_date, input_first_date] if x is not None)
            ld_min = min(x for x in [asset_last_date, input_last_date] if x is not None)
            if helpers.Date.get_difference_in_months(ld_min, fd_max).n < 2:
                raise ShortPeriodLengthError(
                    f"{asset_item.symbol} historical data period length is too short. It must be at least 3 months."
                )

            # append data to dictionaries
            currencies[asset_item.symbol] = asset_item.currency
            names[asset_item.symbol] = asset_item.name
            first_dates[asset_item.symbol] = asset_first_date
            last_dates[asset_item.symbol] = asset_last_date
            own_first_dates[asset_item.symbol] = asset_own_first_date
            own_last_dates[asset_item.symbol] = asset_own_last_date

        # Concatenate all series at once (more efficient than repeated pd.concat in loop)
        df = pd.concat(ror_series_list, axis=1, join="inner")

        first_dates[base_currency_ticker] = currency_first_date
        last_dates[base_currency_ticker] = currency_last_date
        own_last_dates[base_currency_ticker] = currency_last_date
        own_first_dates[base_currency_ticker] = currency_first_date
        currencies["asset list"] = base_currency_ticker

        # get first and last dates
        first_date_list = list(first_dates.values()) + [input_first_date]
        last_date_list = list(last_dates.values()) + [input_last_date]
        list_first_date = max(x for x in first_date_list if x is not None)
        list_last_date = min(x for x in last_date_list if x is not None)

        # range of last and first dates not limited by AssetList first_date & last_date parameters
        own_first_dates_sorted: list = sorted(own_first_dates.items(), key=lambda y: y[1])
        own_last_dates_sorted: list = sorted(own_last_dates.items(), key=lambda y: y[1])

        if isinstance(df, pd.Series):
            # required to convert Series to DataFrame for single asset list
            df = df.to_frame()
        df.columns.name = "Symbols"  # required for Plotly charts

        return dict(
            first_date=list_first_date,
            last_date=list_last_date,
            newest_asset=own_first_dates_sorted[-1][0],
            eldest_asset=own_first_dates_sorted[0][0],
            names_dict=names,
            currencies_dict=currencies,
            own_first_dates_sorted=dict(own_first_dates_sorted),
            own_last_dates_sorted=dict(own_last_dates_sorted),
            ror=df,
        )

    def _make_ror(self, list_asset: asset.Asset, base_currency_name: str) -> pd.Series:
        """
        Make asset rate of return time series.

        Parameters
        ----------
        list_asset : asset.Asset
            Asset object.
        base_currency_name : str
            Name of the base currency.

        Returns
        -------
        pd.Series
            Rate of return monthly time series for the asset.
        """
        asset_currency_name = list_asset.currency
        if asset_currency_name == base_currency_name:
            ror = list_asset.ror
        else:
            asset_currency = asset.Asset(
                symbol=f"{asset_currency_name}{base_currency_name}.FX",
                first_date=list_asset.first_date,
                last_date=list_asset.last_date,
            )
            ror = self._adjust_ror_to_currency(returns=list_asset.ror, asset_currency=asset_currency)
        return ror

    @classmethod
    def _adjust_ror_to_currency(cls, returns: pd.Series, asset_currency: asset.Asset) -> pd.Series:
        """
        Adjust returns time series to a certain currency.

        Parameters
        ----------
        returns : pd.Series
            Rate of return time series.
        asset_currency : asset.Asset
            Currency asset object.

        Returns
        -------
        pd.Series
            Adjusted rate of return time series.
        """
        asset_mult = returns + 1.0
        currency_mult = asset_currency.ror + 1.0
        # join dataframes to have the same Time Series Index
        df = pd.concat([asset_mult, currency_mult], axis=1, join="inner")
        currency_mult = df.iloc[:, -1]
        asset_mult = df.iloc[:, 0]
        x = asset_mult * currency_mult - 1.0
        x.rename(returns.name, inplace=True)
        return x

    def _adjust_price_to_currency_monthly(self, price: pd.Series, asset_currency: str) -> pd.Series:
        """
        Adjust monthly time series of dividends or close values to a base currency.

        Parameters
        ----------
        price : pd.Series
            Price time series.
        asset_currency : str
            Currency of the asset.

        Returns
        -------
        pd.Series
            Adjusted price time series.
        """
        ccy_symbol = f"{asset_currency}{self.currency}.FX"
        currency_rate = asset.Asset(
            ccy_symbol, first_date=self.first_date, last_date=self.last_date
        ).close_monthly.to_frame()
        merged = price.to_frame().join(currency_rate, how="left")
        if merged.isnull().values.any():
            # can happen if the first value is missing
            merged.fillna(method="backfill", inplace=True)
        return merged.iloc[:, 0].multiply(merged[ccy_symbol], axis=0)

    @staticmethod
    def _define_symbol_list(assets: List[Union[str, Type]]) -> List[str]:
        """
        Extract symbols from a list of assets.

        Parameters
        ----------
        assets : list
            List of assets (can be strings or Asset-like objects).

        Returns
        -------
        list of str
            List of symbols.
        """
        return [asset.symbol if hasattr(asset, "symbol") else asset for asset in assets]

    def _add_inflation(self) -> pd.DataFrame:
        """
        Add inflation column to returns DataFrame.

        Returns
        -------
        pd.DataFrame
            DataFrame with assets' rate of return and inflation.
        """
        if hasattr(self, "inflation"):
            return pd.concat([self._assets_ror, self.inflation_ts], axis=1, join="inner")
        else:
            return self._assets_ror

    def _validate_period(self, period: Any) -> None:
        """
        Validate trailing period.

        Parameters
        ----------
        period : int
            Trailing period in years.

        Raises
        ------
        ValueError
            If `period` is not positive or exceeds the available history.
        TypeError
            If `period` is not an integer.
        """
        validators.validate_integer("period", period, min_value=0, inclusive=False)
        if period > self.pl.years:
            raise ValueError(f"'period' ({period}) is beyond historical data range ({self.period_length}).")

    def _get_single_asset_dividends(self, tick: str, remove_forecast: bool = True) -> pd.Series:
        """
        Get monthly dividend time series for a single symbol and adjust to the currency.

        Parameters
        ----------
        tick : str
            Asset symbol.
        remove_forecast : bool, default True
            If True, remove forecasted (future) data.

        Returns
        -------
        pd.Series
            Monthly dividend time series.
        """
        asset = self.asset_obj_dict[tick]
        s = asset.dividends[self.first_date : self.last_date]
        if asset.currency != self.currency:
            s = self._adjust_price_to_currency_monthly(s, asset.currency)
        if remove_forecast:
            s = s[: pd.Period.now(freq="M")]  # Period.now() must be without arguments to be compatible with pandas 2.0
        # Create time series with zeros to pad the empty spaces in dividends time series
        index = pd.date_range(start=self.first_date, end=self.last_date, freq="MS")  # 'MS' to include the last period
        period = index.to_period("M")
        pad_s = pd.Series(data=0, index=period)
        return s.add(pad_s, fill_value=0)

    def _get_assets_dividends(self, remove_forecast=True) -> pd.DataFrame:
        """
        Get monthly dividend time series for all assets. Dividend values are adjusted to base currency.

        Parameters
        ----------
        remove_forecast : bool, default True
            If True, all forecasted (future) data is removed from the time series.

        Returns
        -------
        pd.DataFrame
            Monthly dividend time series for all assets.
        """
        if self._assets_dividends_ts.empty:
            dic = {}
            for tick in self.symbols:
                s = self._get_single_asset_dividends(tick, remove_forecast=remove_forecast)
                dic[tick] = s
            self._assets_dividends_ts = pd.DataFrame(dic)
        return self._assets_dividends_ts

    def _make_real_return_time_series(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Calculate real monthly return time series.

        Rate of return monthly data is adjusted for inflation.

        Parameters
        ----------
        df : pd.DataFrame
            DataFrame of monthly returns.

        Returns
        -------
        pd.DataFrame
            DataFrame of real monthly returns.

        Raises
        ------
        ValueError
            If inflation data is not available.
        """
        if not hasattr(self, "inflation"):
            raise ValueError("Real Return is not defined. Set inflation=True when initiating the class.")
        df = (1.0 + df).divide(1.0 + self.inflation_ts, axis=0) - 1.0
        df.drop(columns=[self.inflation], inplace=True)
        return df

    @property
    def _assets_dividend_yield(self) -> pd.DataFrame:
        """
        Calculate last twelve months (LTM) dividend yield time series (monthly) for each asset.

        LTM dividend yield is the sum trailing twelve months of common dividends per share divided by
        the current price per share.

        All yields are calculated in the asset list base currency after adjusting the dividends and price time series.
        Forecasted (future) dividends are removed.
        Zero value time series are created for assets without dividends.

        Returns
        -------
        pd.DataFrame
            Monthly time series of LTM dividend yield for each asset.
        """
        if self._dividend_yield.empty:
            frame = {}
            df = self._get_assets_dividends(remove_forecast=True)
            for tick in self.symbols:
                div_monthly = df[tick]
                if div_monthly.sum() != 0:
                    asset = self.asset_obj_dict[tick]
                    price_monthly_ts = asset.close_monthly.loc[self.first_date : self.last_date]
                    if asset.currency != self.currency:
                        price_monthly_ts = self._adjust_price_to_currency_monthly(price_monthly_ts, asset.currency)
                else:
                    # skipping prices if no dividends
                    div_yield = div_monthly
                    frame[tick] = div_yield
                    continue
                # Get dividend yield time series
                div_yield = pd.Series(dtype=float)
                div_monthly.index = div_monthly.index.to_timestamp(how="start")
                for date in price_monthly_ts.index.to_timestamp(how="End"):
                    date0 = date - pd.DateOffset(months=12)  # last 12 months
                    ltm_div = div_monthly[date0:date].sum()
                    last_price = price_monthly_ts.loc[:date].iloc[-1]
                    value = ltm_div / last_price
                    div_yield.at[date] = value
                div_yield.index = div_yield.index.to_period("M")
                frame[tick] = div_yield
            self._dividend_yield = pd.DataFrame(frame)
        return self._dividend_yield

    @property
    def _list_of_asset_like_objects(self) -> List[Union[str, Type]]:
        """
        Return list which may include tickers or asset like objects (Portfolio, Asset).

        Returns
        -------
        list
        """
        assets = self._assets
        if not isinstance(assets, list):
            raise ValueError("Assets must be a list.")
        return assets

    @property
    def assets_ror(self) -> pd.DataFrame:
        """
        Rate of return monthly time series for all assets.

        Returns
        -------
        DataFrame
            Rate of return monthly data.
        """
        return self._assets_ror

    @property
    def symbols(self) -> List[str]:
        """
        Return a list of used financial symbols.

        Symbols are similar to tickers but have a namespace information:

        * SPY.US is a symbol
        * SPY is a ticker

        Returns
        -------
        list of str
            List of used symbols.
        """
        return self._define_symbol_list(self._list_of_asset_like_objects)

    @property
    def tickers(self) -> List[str]:
        """
        Return a list of used tickers (symbols without a namespace).

        Tickers are similar to symbols but do not have namespace information:

        * SPY is a ticker
        * SPY.US is a symbol

        Returns
        -------
        list of str
            List of used tickers.
        """
        return [x.split(".", 1)[0] for x in self.symbols]

    @property
    def currency(self) -> str:
        """
        Return the base currency.

        Such properties as rate of return and risk are adjusted to the base currency.

        Returns
        -------
        str
            Base currency.
        """
        return self._currency.currency

    def plot_assets(
        self,
        kind: str = "mean",
        tickers: Union[str, list] = "tickers",
        pct_values: bool = False,
        xy_text: tuple = (0, 10),
        **kwargs,
    ) -> Axes:
        """
        Plot asset points on the risk-return chart with annotations.

        Annualized values for risk and return are used.
        Risk is the standard deviation of monthly rate of return time series.
        Return can be an annualized mean return (expected return) or CAGR (compound annual growth rate).

        Parameters
        ----------
        kind : {'mean', 'cagr'}, default 'mean'
            Type of return: annualized mean return (expected return) or CAGR (compound annual growth rate).

        tickers : {'tickers', 'names'} or list of str, default 'tickers'
            Annotation type for assets.

            - 'tickers': asset symbols are shown in the form 'SPY.US'.
            - 'names': asset names are shown (for example, 'SPDR S&P 500 ETF Trust').
            - list of str: custom annotations for each asset.

        pct_values : bool, default False
            If True, show risk and return values in percent.
            If False, show values as decimals.

        xy_text : tuple, default (0, 10)
            The shift of the annotation text (x, y) from the point.

        **kwargs
            Arbitrary keyword arguments passed to matplotlib.pyplot.scatter.

        Returns
        -------
        Axes
            Matplotlib axes object.

        Examples
        --------
        >>> import matplotlib.pyplot as plt
        >>> al = ok.AssetList(['SPY.US', 'AGG.US'], ccy='USD', inflation=False)
        >>> al.plot_assets()
        >>> plt.show()

        Plotting with default parameters values shows expected return, ticker annotations and algebraic values
        for risk and return.
        To use CAGR instead of expected return use kind='cagr'.

        >>> al.plot_assets(kind='cagr',
        ...               tickers=['US Stocks', 'US Bonds'],  # use custom annotations for the assets
        ...               pct_values=True  # risk and return values are in percents
        ...               )
        >>> plt.show()
        """
        risk_monthly = self._assets_ror.std()
        mean_return_monthly = self._assets_ror.mean()
        risks = helpers.Float.annualize_risk(risk_monthly, mean_return_monthly)
        if kind == "mean":
            returns = helpers.Float.annualize_return(self._assets_ror.mean())
        elif kind == "cagr":
            returns = helpers.Frame.get_cagr(self._assets_ror).loc[self.symbols]
        else:
            raise ValueError('kind should be "mean" or "cagr".')
        # set lists for single point scatter
        if len(self.symbols) < 2:
            risks = [risks]
            returns = [returns]
        # set the plot
        ax = plt.gca()
        plt.autoscale(enable=True, axis="year", tight=False)
        ax.margins(0.05, 0.1)  # increase margins on Y-axis from 5% to 10% as `annotate` moves text upwards
        m = 100 if pct_values else 1
        ax.scatter(risks * m, returns * m, zorder=10, **(kwargs or {}))
        # Set the labels
        if tickers == "tickers":
            asset_labels = self.symbols
        elif tickers == "names":
            asset_labels = list(self.names.values())
        else:
            if not isinstance(tickers, list):
                raise ValueError("tickers parameter should be a list of string labels.")
            if len(tickers) != len(self.symbols):
                raise ValueError("labels and tickers must be of the same length")
            asset_labels = tickers
        # draw the points and print the labels
        for label, x, y in zip(asset_labels, risks, returns):
            ax.annotate(
                label,  # this is the text
                (x * m, y * m),  # this is the point to label
                textcoords="offset points",  # how to position the text
                xytext=(xy_text),  # distance from text to points (x,y)
                ha="center",  # horizontal alignment can be left, right or center
            )
        return ax
