import datetime as dt
from collections.abc import Sequence
from typing import Any, Literal, cast, overload

import polars as pl

from bayesline.api._src.equity.calendar import AsyncCalendarLoaderApi
from bayesline.api._src.equity.exposure import AsyncExposureLoaderApi
from bayesline.api._src.equity.exposure_settings import ExposureSettings
from bayesline.api._src.equity.modelconstruction import (
    AsyncFactorModelConstructionLoaderApi,
)
from bayesline.api._src.equity.modelconstruction_settings import (
    ModelConstructionSettings,
)
from bayesline.api._src.equity.report.api_loader import AsyncReportLoaderApi
from bayesline.api._src.equity.report.exposure import (
    AsyncExposureReportAccessor,
    ExposureReportSettings,
)
from bayesline.api._src.equity.report.factor_stats import (
    AsyncFactorStatsReportAccessor,
    FactorStatsReportSettings,
)
from bayesline.api._src.equity.report.idio import (
    AsyncIdioReportAccessor,
    IdioReportSettings,
)
from bayesline.api._src.equity.report.mcap_exposure import (
    AsyncMcapExposureReportAccessor,
    McapExposureReportSettings,
)
from bayesline.api._src.equity.report.universe import (
    AsyncUniverseCountReportAccessor,
    UniverseCountReportSettings,
)
from bayesline.api._src.equity.report.weights import (
    AsyncWeightsReportAccessor,
    WeightsReportSettings,
)
from bayesline.api._src.equity.riskmodels import (
    AsyncFactorModelApi,
    AsyncFactorModelEngineApi,
    AsyncFactorModelLoaderApi,
    AsyncFactorModelMetadataApi,
)
from bayesline.api._src.equity.riskmodels_settings import (
    FactorRiskModelMetadata,
    FactorRiskModelSettings,
    FactorRiskModelSettingsMenu,
    GetModelMode,
)
from bayesline.api._src.equity.universe import AsyncUniverseLoaderApi
from bayesline.api._src.equity.universe_settings import UniverseSettings
from bayesline.api._src.registry import AsyncSettingsRegistry
from bayesline.api._src.tasks import AsyncTask
from bayesline.api._src.types import DateLike, IdType


class AsyncFactorModelImpl(AsyncFactorModelApi):  # noqa: D101

    def __init__(  # noqa: D102
        self,
        settings: FactorStatsReportSettings,
        report_loader: AsyncReportLoaderApi,
    ) -> None:
        self._settings = settings

        self._universe_settings = [
            cast(UniverseSettings, u)
            for u in self._settings.factor_model_settings.universe
        ]
        self._exposure_settings = [
            cast(ExposureSettings, e)
            for e in self._settings.factor_model_settings.exposures
        ]
        self._modelconstruction_settings = [
            cast(ModelConstructionSettings, mc)
            for mc in self._settings.factor_model_settings.modelconstruction
        ]

        self._report_loader = report_loader

        self._universe_count_report_accessor: (
            AsyncUniverseCountReportAccessor | None
        ) = None
        self._estimation_universe_count_report_accessor: (
            AsyncUniverseCountReportAccessor | None
        ) = None
        self._factor_stats_report_accessor: AsyncFactorStatsReportAccessor | None = None
        self._exposures_report_accessor: dict[int, AsyncExposureReportAccessor] = {}
        self._mcap_exposure_report_accessor: AsyncMcapExposureReportAccessor | None = (
            None
        )
        self._idio_report_accessor: AsyncIdioReportAccessor | None = None
        self._weights_report_accessor: AsyncWeightsReportAccessor | None = None

    @property
    def settings(self) -> FactorRiskModelSettings:  # noqa: D102
        return self._settings.factor_model_settings

    async def _get_universe_count_report_accessor(
        self,
    ) -> AsyncUniverseCountReportAccessor:
        if not self._universe_count_report_accessor:
            universe_count_report_settings = UniverseCountReportSettings(
                universe_settings=self._universe_settings[0],
            )
            universe_count_report_api = await self._report_loader.load(
                universe_count_report_settings
            )
            self._universe_count_report_accessor = (
                await universe_count_report_api.calculate(
                    start_date=None, end_date=None
                )
            )
        return self._universe_count_report_accessor

    async def _get_estimation_universe_count_report_accessor(
        self,
    ) -> AsyncUniverseCountReportAccessor:
        mcs = self._modelconstruction_settings[0]
        estimation_universe_settings: UniverseSettings | None = None
        if mcs.estimation_universe is None:
            return await self._get_universe_count_report_accessor()
        else:
            estimation_universe_settings = cast(
                UniverseSettings, mcs.estimation_universe
            )

        universe_settings = self._universe_settings[0]
        intersection_universe_settings = universe_settings.model_copy(
            update={
                "categorical_filters": [
                    *universe_settings.categorical_filters,
                    *estimation_universe_settings.categorical_filters,
                ],
            }
        )
        if not self._estimation_universe_count_report_accessor:
            estimation_universe_count_report_settings = UniverseCountReportSettings(
                universe_settings=intersection_universe_settings
            )
            estimation_universe_count_report_api = await self._report_loader.load(
                estimation_universe_count_report_settings
            )
            self._estimation_universe_count_report_accessor = (
                await estimation_universe_count_report_api.calculate(
                    start_date=None, end_date=None
                )
            )
        return self._estimation_universe_count_report_accessor

    async def _get_factor_stats_report_accessor(self) -> AsyncFactorStatsReportAccessor:
        if not self._factor_stats_report_accessor:
            report_api = await self._report_loader.load(self._settings)
            self._factor_stats_report_accessor = await report_api.calculate(
                start_date=None, end_date=None
            )
        return self._factor_stats_report_accessor

    async def _get_exposures_report_accessor(
        self, stage: int
    ) -> AsyncExposureReportAccessor:
        if stage not in self._exposures_report_accessor:
            universe_stage = min(stage, 1)  # always use the first universe for now
            mcs_stage = stage if len(self._modelconstruction_settings) > 1 else 1
            mcs_settings = self._modelconstruction_settings[mcs_stage - 1]

            std_universe: UniverseSettings | None = None
            if mcs_settings.estimation_universe is not None:
                std_universe = cast(UniverseSettings, mcs_settings.estimation_universe)

            exposure_report_settings = ExposureReportSettings(
                universe_settings=self._universe_settings[universe_stage - 1],
                exposure_settings=self._exposure_settings[stage - 1],
                standardize_universe_settings=std_universe,
            )
            exposure_report_api = await self._report_loader.load(
                exposure_report_settings
            )
            self._exposures_report_accessor[stage] = (
                await exposure_report_api.calculate(start_date=None, end_date=None)
            )
        return self._exposures_report_accessor[stage]

    async def _get_mcap_exposure_report_accessor(
        self,
    ) -> AsyncMcapExposureReportAccessor:
        if not self._mcap_exposure_report_accessor:
            mcap_exposure_report_settings = McapExposureReportSettings(
                factor_model_settings=self._settings.factor_model_settings
            )
            mcap_exposure_report_api = await self._report_loader.load(
                mcap_exposure_report_settings
            )
            self._mcap_exposure_report_accessor = (
                await mcap_exposure_report_api.calculate(start_date=None, end_date=None)
            )
        return self._mcap_exposure_report_accessor

    async def _get_idio_report_accessor(self) -> AsyncIdioReportAccessor:
        if not self._idio_report_accessor:
            idio_report_settings = IdioReportSettings(
                factor_model_settings=self._settings.factor_model_settings.model_copy(
                    update={
                        "exposures": [
                            ExposureSettings(exposures=[])
                            for _ in self._settings.factor_model_settings.universe
                        ],
                        "modelconstruction": [
                            ModelConstructionSettings(return_clip_bounds=(None, None))
                            for _ in self._settings.factor_model_settings.universe
                        ],
                    }
                )
            )
            idio_report_api = await self._report_loader.load(idio_report_settings)
            self._idio_report_accessor = await idio_report_api.calculate(
                start_date=None, end_date=None
            )
        return self._idio_report_accessor

    async def dates(self) -> list[dt.date]:  # noqa: D102
        return await (await self._get_factor_stats_report_accessor()).get_dates()

    async def factors(self) -> dict[str, list[str]]:  # noqa: D102
        return await (await self._get_factor_stats_report_accessor()).get_factors()

    async def exposures(  # noqa: D102
        self,
        *,
        start: DateLike | None = None,
        end: DateLike | None = None,
        id_type: IdType | None = None,
        stage: int = 1,
    ) -> pl.DataFrame:
        if id_type is not None and id_type != "bayesid":
            raise NotImplementedError("only default id type is supported for now")
        report_accessor = await self._get_exposures_report_accessor(stage)
        df = await report_accessor.get(start_date=start, end_date=end)
        return df.rename(
            lambda c: (
                c.replace("^", ".")
                if "^" in c
                else ("bayesid" if c == "asset_id" else c)
            )
        )

    async def universe(  # noqa: D102
        self,
        *,
        start: DateLike | None = None,
        end: DateLike | None = None,
        id_type: IdType | None = None,
        stage: int = 1,
    ) -> pl.DataFrame:
        # stage is not used because all universes have to be the same (for now)
        if stage != 1:
            raise NotImplementedError("only stage 1 is supported for now")
        if id_type is not None and id_type != "bayesid":
            raise NotImplementedError("only default id type is supported for now")
        filters: list[tuple[str, str, Any] | Sequence[tuple[str, str, Any]]] = []
        if start:
            filters.append(("date", ">=", start))
        if end:
            filters.append(("date", "<=", end))
        return (
            (
                await (
                    await self._get_universe_count_report_accessor()
                ).accessor.get_data(
                    [],
                    expand=("date",),
                    pivot_cols=("asset_id",),
                    value_cols=("count",),
                    filters=filters,
                )
            )
            .select("date", pl.exclude("date").cast(pl.Boolean))
            .rename(lambda x: x.removesuffix("^count"))
        )

    async def estimation_universe(  # noqa: D102
        self,
        *,
        start: DateLike | None = None,
        end: DateLike | None = None,
        id_type: IdType | None = None,
        stage: int = 1,
    ) -> pl.DataFrame:
        # stage is not used because all universes have to be the same (for now)
        if stage != 1:
            raise NotImplementedError("only stage 1 is supported for now")
        if id_type is not None and id_type != "bayesid":
            raise NotImplementedError("only default id type is supported for now")
        filters: list[tuple[str, str, Any] | Sequence[tuple[str, str, Any]]] = []
        if start:
            filters.append(("date", ">=", start))
        if end:
            filters.append(("date", "<=", end))
        return (
            (
                await (
                    await self._get_estimation_universe_count_report_accessor()
                ).accessor.get_data(
                    [],
                    expand=("date",),
                    pivot_cols=("asset_id",),
                    value_cols=("count",),
                    filters=filters,
                )
            )
            .select("date", pl.exclude("date").cast(pl.Boolean))
            .rename(lambda x: x.removesuffix("^count"))
        )

    async def market_caps(  # noqa: D102
        self,
        *,
        start: DateLike | None = None,
        end: DateLike | None = None,
        id_type: IdType | None = None,
    ) -> pl.DataFrame:
        if start is not None or end is not None:
            raise NotImplementedError("date filtering is not supported for now")
        if id_type is not None and id_type != "bayesid":
            raise NotImplementedError("only default id type is supported for now")
        df = await (
            await self._get_mcap_exposure_report_accessor()
        ).get_mcap_weighted_exposures()
        return df.rename(lambda c: c.replace("^", ".") if "^" in c else c)

    async def _get_weights_report_accessor(self) -> AsyncWeightsReportAccessor:
        if not self._weights_report_accessor:
            weights_report_settings = WeightsReportSettings(
                universe_settings=self._universe_settings[0],
                weights=self._modelconstruction_settings[0].weights,
            )
            weights_report_api = await self._report_loader.load(weights_report_settings)
            self._weights_report_accessor = await weights_report_api.calculate(
                start_date=None, end_date=None
            )
        return self._weights_report_accessor

    async def weights(  # noqa: D102
        self,
        *,
        start: DateLike | None = None,
        end: DateLike | None = None,
        id_type: IdType | None = None,
    ) -> pl.DataFrame:
        if id_type is not None and id_type != "bayesid":
            raise NotImplementedError("only default id type is supported for now")
        filters: list[tuple[str, str, Any] | Sequence[tuple[str, str, Any]]] = []
        if start:
            filters.append(("date", ">=", start))
        if end:
            filters.append(("date", "<=", end))
        return (
            await (await self._get_weights_report_accessor()).accessor.get_data(
                [],
                expand=("date",),
                pivot_cols=("asset_id",),
                value_cols=("weight",),
                filters=filters,
            )
        ).rename(lambda x: x.removesuffix("^weight"))

    async def future_asset_returns(  # noqa: D102
        self,
        *,
        start: DateLike | None = None,
        end: DateLike | None = None,
        id_type: IdType | None = None,
    ) -> pl.DataFrame:
        if id_type is not None and id_type != "bayesid":
            raise NotImplementedError("only default id type is supported for now")
        filters: list[tuple[str, str, Any] | Sequence[tuple[str, str, Any]]] = []
        if start:
            filters.append(("date", ">=", start))
        if end:
            filters.append(("date", "<=", end))
        accessor = await self._get_idio_report_accessor()
        # Shift returns back by one date to match V1 forward-return convention:
        # V1 labels the return from t to t+1 at date t; the report labels it at t+1.
        return (
            (
                await accessor.accessor.get_data(
                    [],
                    expand=("date",),
                    pivot_cols=("asset_id",),
                    value_cols=("idio_return",),
                    filters=filters,
                )
            )
            .rename(lambda x: x.removesuffix("^idio_return"))
            .with_columns(pl.exclude("date").shift(-1))
        )

    async def fret(  # noqa: D102
        self,
        *,
        freq: str | None = None,
        cumulative: bool = False,
        start: DateLike | None = None,
        end: DateLike | None = None,
    ) -> pl.DataFrame:
        if freq is not None:
            raise NotImplementedError("only daily frequency is supported for now")
        if cumulative:
            raise NotImplementedError("cumulative returns are not supported for now")
        if start is not None or end is not None:
            raise NotImplementedError("date filtering is not supported for now")
        df = await (await self._get_factor_stats_report_accessor()).get_factor_returns()
        return df.rename(lambda c: c.replace("^", ".") if "^" in c else c)

    async def t_stats(self) -> pl.DataFrame:  # noqa: D102
        df = await (await self._get_factor_stats_report_accessor()).get_t_stats()
        return df.rename(lambda c: c.replace("^", ".") if "^" in c else c)

    async def p_values(self) -> pl.DataFrame:  # noqa: D102
        df = await (await self._get_factor_stats_report_accessor()).get_p_values()
        return df.rename(lambda c: c.replace("^", ".") if "^" in c else c)

    async def r2(self) -> pl.DataFrame:  # noqa: D102
        raise NotImplementedError()

    async def sigma2(self) -> pl.DataFrame:  # noqa: D102
        raise NotImplementedError()


class AsyncFactorModelEngineImpl(AsyncFactorModelEngineApi):  # noqa: D101

    def __init__(  # noqa: D102
        self,
        report_settings: FactorStatsReportSettings,
        report_loader: AsyncReportLoaderApi,
    ) -> None:
        self._report_settings = report_settings
        self._report_loader = report_loader

    @property
    def settings(self) -> FactorRiskModelSettings:  # noqa: D102
        return self._report_settings.factor_model_settings

    async def get_model(  # noqa: D102
        self, mode: GetModelMode = "get-or-compute-and-persist"
    ) -> AsyncFactorModelApi:
        return AsyncFactorModelImpl(self._report_settings, self._report_loader)

    async def get_model_as_task(  # noqa: D102
        self,
        mode: GetModelMode = "get-or-compute-and-persist",
    ) -> AsyncTask[AsyncFactorModelApi]:
        raise NotImplementedError()


class AsyncFactorModelLoaderImpl(AsyncFactorModelLoaderApi):  # noqa: D101

    def __init__(  # noqa: D102
        self,
        *,
        settings_registry: AsyncSettingsRegistry[
            FactorRiskModelSettings,
            FactorRiskModelSettingsMenu,
        ],
        calendar_api: AsyncCalendarLoaderApi,
        universe_api: AsyncUniverseLoaderApi,
        exposure_api: AsyncExposureLoaderApi,
        modelconstruction_api: AsyncFactorModelConstructionLoaderApi,
        report_loader: AsyncReportLoaderApi,
        metadata_api: AsyncFactorModelMetadataApi,
    ) -> None:
        self._settings_registry = settings_registry
        self._calendar_api = calendar_api
        self._universe_api = universe_api
        self._exposure_api = exposure_api
        self._modelconstruction_api = modelconstruction_api
        self._report_loader = report_loader
        self._metadata_api = metadata_api

    @property
    def settings(  # noqa: D102
        self,
    ) -> AsyncSettingsRegistry[FactorRiskModelSettings, FactorRiskModelSettingsMenu]:
        return self._settings_registry

    @overload
    async def list_riskmodels(  # noqa: D102
        self, risk_dataset: str | None = None, *, df: Literal[False] = ...
    ) -> list[FactorRiskModelMetadata]: ...

    @overload
    async def list_riskmodels(  # noqa: D102
        self, risk_dataset: str | None = None, *, df: Literal[True] = ...
    ) -> pl.DataFrame: ...

    async def list_riskmodels(  # noqa: D102
        self, risk_dataset: str | None = None, *, df: bool = False
    ) -> list[FactorRiskModelMetadata] | pl.DataFrame:
        if df:
            return await self._metadata_api.list_riskmodels(risk_dataset, df=True)
        return await self._metadata_api.list_riskmodels(risk_dataset, df=False)

    async def load(  # noqa: D102
        self, ref_or_settings: str | int | FactorRiskModelSettings
    ) -> AsyncFactorModelEngineApi:
        if isinstance(ref_or_settings, FactorRiskModelSettings):
            settings = ref_or_settings
        else:
            settings = await self._settings_registry.get(ref_or_settings)

        resolved_universe = [
            (
                u
                if isinstance(u, UniverseSettings)
                else await self._universe_api.settings.get(u)
            )
            for u in settings.universe
        ]

        datasets = set([u.dataset for u in resolved_universe if u.dataset is not None])

        if len(datasets) == 0:
            raise ValueError("Must define a dataset inside the UniverseSettings")
        elif len(datasets) == 1:
            dataset = datasets.pop()
        else:
            raise ValueError(
                "datasets for multi pass risk model must be the same, "
                f"but got {datasets}"
            )

        resolved_universe = [
            u.model_copy(update={"dataset": dataset}) for u in resolved_universe
        ]

        resolved_exposures = [
            (
                e
                if isinstance(e, ExposureSettings)
                else await self._exposure_api.settings.get(e)
            )
            for e in settings.exposures
        ]

        async def resolve_est_universe(
            mc: ModelConstructionSettings,
        ) -> ModelConstructionSettings:
            if isinstance(mc.estimation_universe, str | int):
                est_universe_settings = await self._universe_api.settings.get(
                    mc.estimation_universe
                )
                return mc.model_copy(
                    update={"estimation_universe": est_universe_settings}
                )
            elif isinstance(mc.estimation_universe, UniverseSettings):
                return mc
            elif mc.estimation_universe is None:
                return mc
            else:
                raise ValueError(
                    f"Invalid estimation universe: {mc.estimation_universe}"
                )

        resolved_modelconstruction = [
            (
                await resolve_est_universe(mc)
                if isinstance(mc, ModelConstructionSettings)
                else await self._modelconstruction_api.settings.get(mc)
            )
            for mc in settings.modelconstruction
        ]

        settings_resolved = settings.model_copy(
            update={
                "universe": resolved_universe,
                "exposures": resolved_exposures,
                "modelconstruction": resolved_modelconstruction,
            },
        )

        report_settings = FactorStatsReportSettings(
            factor_model_settings=settings_resolved,
        )

        return AsyncFactorModelEngineImpl(report_settings, self._report_loader)
