import abc
from typing import Any, Literal

import polars as pl
from pydantic import Field

from bayesline.api._src._utils import docstrings_from_sync
from bayesline.api._src.equity.portfolioreport import (
    AsyncReportAccessorApi,
    ReportAccessorApi,
)
from bayesline.api._src.equity.report.accessor import (
    AsyncTypedReportAccessorApi,
    TypedReportAccessorApi,
)
from bayesline.api._src.equity.report.api import AsyncReportApi, ReportApi
from bayesline.api._src.equity.report.settings import (
    FactorCovarianceSettings,
    ReportSettings,
)
from bayesline.api._src.equity.riskmodels_settings import FactorRiskModelSettings
from bayesline.api._src.tasks import AsyncTask, Task
from bayesline.api._src.types import DateLike


class FactorCovarianceReportSettings(ReportSettings):
    """Settings for a factor covariance report."""

    report_type: Literal["factor_covariance"] = Field(
        default="factor_covariance", description="Report type discriminator"
    )
    factor_model_settings: FactorRiskModelSettings = Field()
    factor_cov_settings: FactorCovarianceSettings = Field(
        default_factory=FactorCovarianceSettings,
        description="Settings for the factor covariance matrix forecasts.",
    )


class FactorCovarianceReportAccessor(TypedReportAccessorApi):
    """Specific accessor for a exposure report."""

    @abc.abstractmethod
    def get_covariance(self) -> pl.DataFrame:
        """
        Get the time-series of factor covariances for the factors.

        Returns
        -------
        pl.DataFrame
            A dataframe with a `date` and `factor` columns and columns for each factor
            with their respective factor covariances.
        """

    @abc.abstractmethod
    def get_correlation(self) -> pl.DataFrame:
        """
        Get the time-series of factor correlations for the factors.

        Returns
        -------
        pl.DataFrame
            A dataframe with a `date` and `factor` columns and columns for each factor
            with their respective factor correlations.
        """


@docstrings_from_sync
class AsyncFactorCovarianceReportAccessor(AsyncTypedReportAccessorApi):

    @abc.abstractmethod
    async def get_covariance(self) -> pl.DataFrame: ...  # noqa: D102

    @abc.abstractmethod
    async def get_correlation(self) -> pl.DataFrame: ...  # noqa: D102


class FactorCovarianceReportApi(
    ReportApi[
        [DateLike, DateLike],
        FactorCovarianceReportAccessor,
        FactorCovarianceReportSettings,
    ]
):
    """API for a exposure report."""

    @abc.abstractmethod
    def get_maybe_cached(
        self,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> FactorCovarianceReportAccessor | None:
        """
        Get the cached factor covariance report accessor API if it is available.

        Parameters
        ----------
        start_date: DateLike | None
            The start date of the report.
        end_date: DateLike | None
            The end date of the report.
        **kwargs: Any
            Additional keyword arguments.

        Returns
        -------
        FactorCovarianceReportAccessor | None
            The cached factor covariance report accessor API, or None if the report is
            not cached.
        """

    @abc.abstractmethod
    def calculate(
        self,
        start_date: DateLike,
        end_date: DateLike,
        **kwargs: Any,
    ) -> FactorCovarianceReportAccessor:
        """
        Calculate the factor covariance report.

        Parameters
        ----------
        start_date: DateLike
            The start date of the report.
        end_date: DateLike
            The end date of the report.
        **kwargs: Any
            Additional keyword arguments.

        Returns
        -------
        FactorCovarianceReportAccessor
            The exposure report accessor.
        """

    @abc.abstractmethod
    def calculate_as_task(
        self,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> Task[FactorCovarianceReportAccessor]:
        """
        Calculate the factor covariance report as a task.

        Parameters
        ----------
        start_date: DateLike | None
            The start date of the report.
        end_date: DateLike | None
            The end date of the report.
        **kwargs: Any
            Additional keyword arguments.

        Returns
        -------
        Task[FactorCovarianceReportAccessor]
            The task for the factor covariance report accessor.
        """


@docstrings_from_sync
class AsyncFactorCovarianceReportApi(
    AsyncReportApi[
        [DateLike | None, DateLike | None],
        AsyncFactorCovarianceReportAccessor,
        FactorCovarianceReportSettings,
    ]
):

    @abc.abstractmethod
    async def get_maybe_cached(  # noqa: D102
        self,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> AsyncFactorCovarianceReportAccessor | None: ...

    @abc.abstractmethod
    async def calculate(  # noqa: D102
        self,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> AsyncFactorCovarianceReportAccessor: ...

    @abc.abstractmethod
    async def calculate_as_task(  # noqa: D102
        self,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> AsyncTask[AsyncFactorCovarianceReportAccessor]: ...


class FactorCovarianceReportAccessorImpl(FactorCovarianceReportAccessor):  # noqa: D101

    def __init__(self, accessor: ReportAccessorApi):
        self._accessor = accessor

    @property
    def accessor(self) -> ReportAccessorApi:  # noqa: D102
        return self._accessor

    def get_covariance(self) -> pl.DataFrame:  # noqa: D102
        df = self.accessor.get_data(
            [],
            expand=("date", "factor"),
            value_cols=("covariance",),
            pivot_cols=("factor_col",),
        )
        return df.rename(lambda x: x.removesuffix("^covariance"))

    def get_correlation(self) -> pl.DataFrame:  # noqa: D102
        df = self.accessor.get_data(
            [],
            expand=("date", "factor"),
            value_cols=("correlation",),
            pivot_cols=("factor_col",),
        )
        return df.rename(lambda x: x.removesuffix("^correlation"))


class AsyncFactorCovarianceReportAccessorImpl(  # noqa: D101
    AsyncFactorCovarianceReportAccessor
):

    def __init__(self, accessor: AsyncReportAccessorApi):
        self._accessor = accessor

    @property
    def accessor(self) -> AsyncReportAccessorApi:  # noqa: D102
        return self._accessor

    async def get_covariance(self) -> pl.DataFrame:  # noqa: D102
        df = await self.accessor.get_data(
            [],
            expand=("date", "factor"),
            value_cols=("covariance",),
            pivot_cols=("factor_col",),
        )
        return df.rename(lambda x: x.removesuffix("^covariance"))

    async def get_correlation(self) -> pl.DataFrame:  # noqa: D102
        df = await self.accessor.get_data(
            [],
            expand=("date", "factor"),
            value_cols=("correlation",),
            pivot_cols=("factor_col",),
        )
        return df.rename(lambda x: x.removesuffix("^correlation"))
