# Copyright (c) ZhangYundi.
# Licensed under the MIT License. 
# Created on 2025/12/17 10:28
# Description:
from functools import partial
from typing import Literal, Sequence, Any

import logair
import pandas as pd
import polars as pl
from lidb import DataLoader, Dataset, from_polars
from lidb.qdf import QDF, Expr
from polars import selectors as cs, DataFrame
from tqdm.auto import tqdm

from .report import get_report
from .strategy import Strategy
from .utils import demean_forward_returns

logger = logair.get_logger(__name__)


class Zoo:
    """用于研究alpha因子的数据管理部分"""

    _factor_dls: DataLoader
    _add_dls: DataLoader
    _label_dls: DataLoader
    _mount_dls: DataLoader
    _filter_dls: DataLoader
    _filted_index: DataFrame
    _loader_params: dict[Any, Any]
    _factors: DataFrame
    _qdf: QDF

    def __init__(self):
        self.logger = logair.get_logger(f"{__name__}.{self.__class__.__name__}")
        self.init()

    def init(self):
        """初始化"""
        self._factor_dls = DataLoader("factor", )
        self._add_dls = DataLoader("adding", )
        self._label_dls = DataLoader("label", )
        self._mount_dls = DataLoader("mount", )
        self._filter_dls = DataLoader("filter", )
        self._filted_index: pl.DataFrame = None
        self._loader_params = dict()
        self._factors: pl.DataFrame = None # 过滤后的因子数据
        self._qdf: QDF = None


    def load(self,
             beg_date: str,
             end_date: str,
             times: list[str],
             ds_list: list[Dataset],
             n_jobs: int = -1,
             backend: Literal["threading", "multiprocessing", "loky"] = "loky",
             **constraints):
        self.init()
        self._loader_params = dict(beg_date=beg_date,
                                   end_date=end_date,
                                   times=times,
                                   n_jobs=n_jobs,
                                   backend=backend,
                                   eager=True,
                                   **constraints)
        self._factor_dls.get(ds_list=ds_list, **self._loader_params)
        self._update_data(self._factor_dls.data)
        return self.filter()

    def _update_data(self, df: pl.LazyFrame | pl.DataFrame):
        """更新 factors data"""
        if isinstance(df, pl.LazyFrame):
            df = df.collect()
        index = ("date", "time", "asset")
        if self._factors is None:
            self._factors = df
        else:
            self._factors = self._factors.join(df, on=index, how="left")
        if self._qdf is not None:
            self._qdf.data = self._qdf.data.join(df, on=index, how="left")

    def add_dataset(self, ds_list: list[Dataset]):
        """添加新的数据集"""
        if ds_list:
            self._add_dls.get(ds_list=ds_list, **self._loader_params)
            self._update_data(self._add_dls.data)
        return self

    def to_qdf(self, data: pl.DataFrame | pl.LazyFrame):
        self._qdf = from_polars(data, align=True)
        return self

    def label(self, ds_list: list[Dataset]):
        if ds_list:
            self._label_dls.get(ds_list=ds_list, **self._loader_params)
            self._update_data(self._label_dls.data)
        return self

    def mount(self, ds_list: list[Dataset] = None, columns: list[str] | None = None, avoid_future: bool = True):
        """挂载数据: 比如alphalens需要的行业(group)/收益数据..."""
        if ds_list:
            self._mount_dls.get(ds_list=ds_list, avoid_future=avoid_future, columns=columns, **self._loader_params)
            # 更新 factor_data
            self._update_data(self._mount_dls.data)
        return self

    def filter(self, ds_list: list[Dataset] = None):
        """过滤数据集：每个过滤数据集列名必须命名为 `cond` """
        if ds_list:
            index = ("date", "time", "asset")
            self._filter_dls.get(ds_list=[ds.alias(ds.data_name) for ds in ds_list], **self._loader_params)
            filter_df = self._filter_dls.all_data.select(*index, cond=pl.sum_horizontal(cs.exclude(index)).fill_null(1).fill_nan(1))
            self._filted_index = filter_df.filter(pl.col("cond") < 1).select(*index)
            if self._factors is not None:
                self._factors = self._filted_index.join(self._factors, on=index, how="left")
            else:
                self._factors = self._filted_index.join(self._factor_dls.all_data, on=index, how="left")
        else:
            self._filter_dls = DataLoader("filter", )
            self._filted_index = None
            self._factors = self._factor_dls.all_data  # 恢复成过滤前的状态
            self._qdf = None
        return self

    def sql(self, *exprs: str) -> pl.DataFrame:
        if self._qdf is None:
            self.to_qdf(self._factors)
        return self._qdf.sql(*exprs, show_progress=True)

    def _to_alphalens(self,
                      alpha_name: str,
                      demeaned: bool = False,
                      beg_date: str = None,
                      end_date: str = None,
                      times: list[str] = None,
                      bins: int = 10,
                      reverse: bool = False,
                      grouper: Sequence[str] = None):
        """转成 alphalens 所需要的因子分析格式"""
        _grouper = grouper
        grouper = ["date", "time", *_grouper] if _grouper else ["date", "time"]
        index = ("date", "time", "asset")
        # 处理收益列
        if not beg_date:
            beg_date = self.factors["date"].min()
        if not end_date:
            end_date = self.factors["date"].max()
        if not times:
            times = self.factors["time"].unique().to_list()
        factor_data = self.factors.filter(
            (pl.col("date") >= beg_date),
            (pl.col("date") <= end_date),
            (pl.col("time").is_in(times)))
        factor_data = factor_data.sort(by=["date", "time", "asset"])
        if demeaned:
            factor_data = demean_forward_returns(factor_data)
        factor_data = factor_data.with_columns(
            (factor_data[period]).alias(period)
            for period in self.periods
        )

        cols = [*index, *self.periods, *_grouper] if _grouper else [*index, *self.periods]
        if reverse:
            cols.append((-pl.col(alpha_name)).alias("factor"))
        else:
            cols.append(pl.col(alpha_name).alias("factor"))
        labels = [str(i) for i in range(1, bins + 1)]
        cat = pl.Enum(labels)
        return (
            factor_data
            .select(cols)
            .drop_nulls()
            .with_columns(
                pl.col("factor")
                .qcut(bins,
                      labels=labels,
                      allow_duplicates=True)
                .over(grouper)
                .cast(cat)
                .alias("factor_quantile"))
            .cast({pl.Categorical: pl.Utf8})
            .select(*index, "factor", "factor_quantile", *self.periods)
        )

    def report(self,
               exprs: list[str],
               reverse: bool = False,
               lag: int = 1,
               N: int = 10,
               demeaned: bool = False,
               beg_date: str = None,
               end_date: str = None,
               times: list[str] = None,
               grouper: Sequence[str] = None, ):
        """
        qlib 版本的因子报告

        Parameters
        ----------
        exprs: list[str]
            因子表达式
        reverse : bool
        lag: int
            计算因子自回归以及换手时的滞后期数, 默认上一期:1
        N: int
            分组数量, 默认10
        demeaned : bool
            是否进行去均值处理。如果为 True，则在计算收益时会去掉因子的均值影响。
        beg_date: str
        end_date: str
        times: list[str]
        grouper: Sequence[str]
            自定义分组: 用于分组
        Notes
        -----
        group_neutral 和 demeaned 控制是否对收益做调整, group_neutral优先级更高
            - group_neutral 为 True: 收益-行业均值
            - group_neutral 为 False: 如果demeaned 为 True，则收益 - 全市场均值
        """
        index = ("date", "time", "asset")
        logger.info("Calculating alpha expr.")
        expr_data = self.sql(*exprs)
        alpha_names = expr_data.select(cs.exclude(index)).columns

        beg_date = self.factors["date"].min() if beg_date is None else beg_date
        end_date = self.factors["date"].max() if end_date is None else end_date
        _times = self.factors["time"].unique().to_list()
        _times.sort()
        times = _times if times is None else times
        logger.info(f"{beg_date} -> {end_date}: {times}")

        task_num = len(alpha_names)
        # figs_collect = dict()
        with tqdm(total=task_num, desc="Generating report", leave=False) as bar:
            for alpha_name in alpha_names:
                bar.set_postfix_str(alpha_name)
                logger.info(f"REPORT NAME: {alpha_name}")
                _factor_data = self._to_alphalens(alpha_name,
                                                  demeaned=demeaned,
                                                  beg_date=beg_date,
                                                  end_date=end_date,
                                                  times=times,
                                                  bins=N,
                                                  reverse=reverse,
                                                  grouper=grouper, )
                # figs_collect[alpha_name] = model_performance_graph(_factor_data, lag=lag, N=N, show_notebook=show_notebook)
                fig = get_report(_factor_data, factor_name=alpha_name, lag=lag, N=N, )
                fig.show()
                bar.update()
            bar.close()
        # return figs_collect

    def bt(self,
           exprs: list[str],
           strategy: Strategy,
           period: str = "1d",
           reverse: bool = False,
           beg_date: str = None,
           end_date: str = None,
           times: list[str] = None, ):

        """
        多因子回测

        Parameters
        ----------
        exprs: list[str]
        strategy: Strategy
        period: str
        reverse: bool
        beg_date: str
        end_date: str
        times: list[str]

        Examples
        --------

        >>> for name, strategy_partial in zoo.bt():
        >>>     pos_df = strategy_partial()

        """

        index = ("date", "time", "asset")
        mounted_cols = []
        if self.mounted is not None:
            mounted_cols = self.mounted.select(cs.exclude(index)).columns
        # cols = [*index, *mounted_cols]
        logger.info("Calculating alpha expr.")
        self.sql(*exprs)
        # if not period:
        #     period = self.periods[0]
        exprs_alias = [Expr(expr).alias for expr in exprs]

        cols = [*index, *mounted_cols]

        beg_date = self.factors["date"].min() if beg_date is None else beg_date
        end_date = self.factors["date"].max() if end_date is None else end_date
        _times = self.factors["time"].unique().to_list()
        _times.sort()
        times = _times if times is None else times
        logger.info(f"{beg_date} -> {end_date}: {times}")
        ref_df = (self.factors
                  .filter((pl.col("date") >= beg_date),
                          (pl.col("date") <= end_date),
                          (pl.col("time").is_in(times)))
                  .select(*index, *mounted_cols, *exprs_alias)
                  # .join(self.labels.select(*index, period), on=index, how="left")
                  )
        strategy.set_params(bt__period=period)

        for alias in exprs_alias:
            df = ref_df.select(*cols, pl.col(alias).alias("score"))
            if reverse:
                df = df.with_columns((-pl.col("score")).alias("score"))
            yield alias, partial(strategy.run, df)

    @property
    def periods(self) -> list[str]:
        res = list()
        if self._label_dls.all_data is None:
            return res
        index = ("date", "time", "asset")
        cols = self._label_dls.all_data.select(cs.exclude(index)).columns
        for col in cols:
            try:
                pd.Timedelta(col)
                res.append(col)
            except ValueError:
                # 解析失败，说明不是时间差
                pass
        return res

    @property
    def factors(self) -> pl.DataFrame | None:
        """过滤后的因子数据"""
        if self._qdf is not None:
            return self._qdf.data
        return self._factors

    @property
    def mounted(self) -> pl.DataFrame | None:
        return self._mount_dls.all_data

    @property
    def filters(self) -> pl.DataFrame | None:
        return self._filter_dls.all_data

    @property
    def labels(self) -> pl.DataFrame | None:
        return self._label_dls.all_data


zoo = Zoo()
