import abc
import datetime as dt
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 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 FactorStatsReportSettings(ReportSettings):
    """Settings for a factor stats report."""

    report_type: Literal["factor_stats"] = Field(
        default="factor_stats", description="Report type discriminator"
    )
    factor_model_settings: FactorRiskModelSettings = Field()


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

    @abc.abstractmethod
    def get_dates(self) -> list[dt.date]:
        """
        Get the dates covered in the factor stats report.

        Returns
        -------
        list[dt.date]
            The dates covered in the factor stats report.
        """

    @abc.abstractmethod
    def get_factors(self) -> dict[str, list[str]]:
        """
        Get the factor groups and their factor covered in this factor stats report.

        Returns
        -------
        dict[str, list[str]]
            Dict where the keys are the factor groups, and the values are the factors
            in that group.
        """

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

        Returns
        -------
        pl.DataFrame
            A dataframe with a `date` column and columns for each factor with their
            respective factor returns, in factor_group^factor format.
        """

    @abc.abstractmethod
    def get_t_stats(self) -> pl.DataFrame:
        """
        Get the time-series of t-statistics for the factors.

        Returns
        -------
        pl.DataFrame
            A dataframe with a `date` column and columns for each factor with their
            respective t-statistics, in factor_group^factor format.
        """

    @abc.abstractmethod
    def get_p_values(self) -> pl.DataFrame:
        """
        Get the time-series of p-values for the factors.

        Returns
        -------
        pl.DataFrame
            A dataframe with a `date` column and columns for each factor with their
            respective p-values, in factor_group^factor format.
        """

    @abc.abstractmethod
    def get_vif(self) -> pl.DataFrame:
        r"""
        Get the time-series of variance inflation factors (VIF) for the factors.

        The VIF measures multicollinearity — how much the variance of a
        regression coefficient is inflated due to correlation with other
        factors. For factor *j*, VIF is defined as:

        .. math::

            \\text{VIF}_j = \\frac{1}{1 - R_j^2}

        where :math:`R_j^2` is the R-squared from regressing factor *j* on all
        other factors. Equivalently,
        :math:`\\text{VIF}_j = (X^\\top W X)_{jj} \\cdot (X^\\top W X)^{-1}_{jj}`.

        A VIF of 1 means no multicollinearity; values above 5--10 suggest that
        the factor's coefficient estimate is unreliable due to near-linear
        dependence on other factors.

        When the model includes linear constraints (e.g., a sum-to-zero
        constraint on industry factor returns), the unconstrained Gram matrix
        :math:`X^\\top W X` is typically singular or near-singular along the
        constrained directions — for instance, the market intercept column is
        a linear combination of the industry dummy columns. In this case VIF
        is computed in the null space of the constraint matrix, measuring
        multicollinearity only among the degrees of freedom that the model is
        actually estimating. Factors whose collinearity is resolved by the
        constraint will show moderate VIF rather than artificially extreme
        values.

        Returns
        -------
        pl.DataFrame
            A dataframe with a `date` column and columns for each factor with
            their respective VIF values, in factor_group^factor format.
        """


@docstrings_from_sync
class AsyncFactorStatsReportAccessor(AsyncTypedReportAccessorApi):

    @abc.abstractmethod
    async def get_dates(self) -> list[dt.date]: ...  # noqa: D102

    @abc.abstractmethod
    async def get_factors(self) -> dict[str, list[str]]: ...  # noqa: D102

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

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

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

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


class FactorStatsReportApi(
    ReportApi[
        [DateLike, DateLike], FactorStatsReportAccessor, FactorStatsReportSettings
    ]
):
    """API for a exposure report."""

    @abc.abstractmethod
    def get_maybe_cached(
        self,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> FactorStatsReportAccessor | None:
        """
        Get the cached factor stats 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
        -------
        FactorStatsReportAccessor | None
            The cached factor stats report accessor API, or None if the report is not cached.
        """

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

        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
        -------
        FactorStatsReportAccessor
            The exposure report accessor.
        """

    @abc.abstractmethod
    def calculate_as_task(
        self,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> Task[FactorStatsReportAccessor]:
        """
        Calculate the factor stats 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[FactorStatsReportAccessor]
            The task for the factor stats report accessor.
        """


@docstrings_from_sync
class AsyncFactorStatsReportApi(
    AsyncReportApi[
        [DateLike | None, DateLike | None],
        AsyncFactorStatsReportAccessor,
        FactorStatsReportSettings,
    ]
):

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

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

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


class FactorStatsReportAccessorImpl(FactorStatsReportAccessor):  # noqa: D101

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

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

    def get_dates(self) -> list[dt.date]:  # noqa: D102
        df = self.accessor.get_data(
            path=[],
            expand=("date",),
            value_cols=(),
        )
        return df["date"].unique().to_list()

    def get_factors(self) -> dict[str, list[str]]:  # noqa: D102
        df = self.accessor.get_data(
            [],
            expand=("factor_group", "factor"),
            value_cols=(),
        )
        result: dict[str, list[str]] = {}
        for (factor_group,), group_df in df.group_by("factor_group"):
            result[factor_group] = (
                group_df["factor"].unique(maintain_order=True).to_list()
            )
        return result

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

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

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

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


class AsyncFactorStatsReportAccessorImpl(AsyncFactorStatsReportAccessor):  # noqa: D101

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

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

    async def get_dates(self) -> list[dt.date]:  # noqa: D102
        df = await self.accessor.get_data(
            path=[],
            expand=("date",),
            value_cols=(),
        )
        return df["date"].unique().to_list()

    async def get_factors(self) -> dict[str, list[str]]:  # noqa: D102
        df = await self.accessor.get_data(
            [],
            expand=("factor_group", "factor"),
            value_cols=(),
        )
        result: dict[str, list[str]] = {}
        for (factor_group,), group_df in df.group_by("factor_group"):
            result[factor_group] = (
                group_df["factor"].unique(maintain_order=True).to_list()
            )
        return result

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

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

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

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