# Copyright (c) ZhangYundi.
# Licensed under the MIT License. 
# Created on 2025/12/22 16:36
# Description:
from collections import defaultdict
from typing import Sequence

import logair
import polars as pl
import polars.selectors as cs
import xcals
import ygo
from tqdm.auto import tqdm

from ..transformer import as_transformer

logger = logair.get_logger(__name__)


@as_transformer
def backtest(weight_df: pl.DataFrame,
             limit: bool = True,
             show_progress: bool = True) -> dict[str, pl.DataFrame | None]:
    """
    高频多因子回测: 某个时点target_weight如果为 None, 则跳过调仓

    Parameters
    ----------
    weight_df: polars.DataFrame
        需要包含字段: target_weight, price(换仓时点价格), close(日收盘), prev_close(昨收价), open(开盘价)
        如果 limit 为 True，则必须包括 limit 字段
    limit: bool
        是否限制某些买卖：比如涨跌停
    show_progress: bool

    Returns
    -------
    dict[str, polars.DataFrame | None]
        每日的持仓以及换仓详细信息
    """
    check_fields = {"target_weight", "price", "close", "prev_close", "open"}
    assert len(check_fields & set(weight_df.columns)) >= len(
        check_fields), f"Missing fields({check_fields - set(weight_df.columns)}) in weight_df."
    if limit:
        assert "limit" in weight_df.columns, "`limit` not in weight_df.columns"

    weight_df = weight_df.cast({pl.Decimal: pl.Float64})
    # 重新分配时间,价格: prev_close - "0", open: "00", close: "25"
    price_df_cols = ["date", "time", "asset", "price"]
    prev_close = (weight_df
                  .group_by("date", "asset")
                  .agg(time=pl.lit("0"),
                       price=pl.col("prev_close").max())
                  .select(price_df_cols))
    open = (weight_df
            .group_by("date", "asset")
            .agg(time=pl.lit("00"),
                 price=pl.col("open").max())
            .select(price_df_cols))
    close = (weight_df
             .group_by("date", "asset")
             .agg(time=pl.lit("25"),
                  price=pl.col("close").max())
             .select(price_df_cols))
    price_dfs = [prev_close, open, weight_df.select("date", "time", "asset", "price"), close]
    price_df = pl.concat(price_dfs).sort("date", "time", "asset")
    price_df = (price_df.with_columns(price_diff=pl.col("price")
                                      .diff(-1)
                                      .over(partition_by=["date", "asset"],
                                            order_by="time"))
                .with_columns(forward_ret=pl.when(pl.col("price") > 0)
                              .then(-pl.col("price_diff") / pl.col("price"))
                              .otherwise(0.0))
                .drop("price_diff"))

    weight_df = price_df.join(weight_df.drop("price"), on=["date", "time", "asset"], how="left")

    weight_df = weight_df.with_columns(_sum_weight=pl.col("target_weight").sum().over("date", "time"))
    weight_df = weight_df.with_columns(target_weight=pl.col("target_weight") / pl.col("_sum_weight")).drop(
        "_sum_weight")

    total_num = weight_df.select("date").n_unique()
    pbar = tqdm(total=total_num,
                desc="Backtesting",
                ascii="⣀⣄⣤⣦⣶⣷⣿",
                leave=False) if show_progress else None
    beg_date = weight_df["date"][0]
    beg_time = weight_df["time"].unique(maintain_order=True).to_list()[2]  # 过滤掉 0, 00
    time_num = weight_df["time"].n_unique() - 3  # 减掉临时增加的3个虚拟时点

    scalar = 0.0

    pos_list = list()
    cur_date = ""
    for (date, time), pos in weight_df.group_by("date", "time", maintain_order=True):
        if time >= "25":
            continue
        if show_progress:
            pbar.set_postfix_str(f"{date}")
        if not pos_list:
            # 第一次建仓:
            if time in {"0", "00", "25"}:
                continue
            pos = pos.with_columns(avail_sell=pl.lit(0.0), beg_weight=pl.lit(0.0))
        else:
            prev_pos = pos_list[-1]
            prev_pos = prev_pos.select("asset", "avail_sell", beg_weight=pl.col("end_weight"))
            pos = (pos
                   .join(prev_pos, on="asset", how="full", coalesce=True)
                   .with_columns(date=pl.col("date").fill_null(date),
                                 time=pl.col("time").fill_null(time),
                                 avail_sell=pl.col("avail_sell").fill_null(0.0),
                                 beg_weight=pl.col("beg_weight").fill_null(0.0))
                   .with_columns(target_weight=pl.col("target_weight").fill_null(pl.col("beg_weight")))
                   # target_weight 为 null 时，认为是不调仓
                   )
        if date == beg_date:
            scalar += 1 / time_num
        else:
            scalar = 1.0
        target_weight = (pl.when(date=beg_date)
                         .then(pl.col("target_weight") / time_num + pl.col("beg_weight"))
                         .otherwise(pl.col("target_weight")))
        avail_sell = (pl.when(date=beg_date)
                      .then(0.0)
                      .when(time=beg_time)
                      .then(pl.col("beg_weight"))
                      .otherwise(pl.col("avail_sell"))
                      )
        pos = pos.with_columns(target_weight=target_weight, avail_sell=avail_sell)

        chg_weight = (pl.col("target_weight") - pl.col("beg_weight")).clip(lower_bound=-pl.col("avail_sell"))
        pos = pos.with_columns(chg_weight=chg_weight)
        # 涨跌停限制
        if limit:
            chg_weight = pl.when(pl.col("limit") > 0).then(0.0).otherwise(pl.col("chg_weight"))
            pos = pos.with_columns(chg_weight=chg_weight)
        # 更新可卖
        avail_sell = pl.col("avail_sell") + pl.col("chg_weight").clip(upper_bound=0.0)
        # 处理买入： sum(买入) <= sum(卖出)
        total_sell = ((pl.col("chg_weight") < 0) * pl.col("chg_weight")).sum()
        total_buy = ((pl.col("chg_weight") > 0) * pl.col("chg_weight")).sum()
        adj_ratio = 1.0 if date == beg_date else total_sell.abs() / total_buy
        adj_buy = adj_ratio * pl.col("chg_weight")
        change_weight = pl.when(pl.col("chg_weight") > 0).then(adj_buy).otherwise(pl.col("chg_weight"))

        pos = (
            pos
            .with_columns(avail_sell=avail_sell)
            .with_columns(change_weight=change_weight)
            .with_columns(real_weight=pl.col("beg_weight") + pl.col("change_weight"), )
            .with_columns(end_weight=pl.col("real_weight") * (1 + pl.col("forward_ret")))
            .with_columns(end_weight=pl.col("end_weight") / pl.col("end_weight").sum() * scalar)
        )

        pos_list.append(pos)
        if show_progress:
            if date != cur_date:
                cur_date = date
                pbar.update()
    if pbar is not None:
        pbar.close()

    res = defaultdict(None)

    if pos_list:
        pos_df = pl.concat(pos_list).sort("date", "time", "asset")
        pos_df = pos_df.filter(pl.col("time") < "25")
        ret_df = (
            pos_df
            .group_by("date", "time")
            .agg(long=(pl.col("real_weight") * pl.col("forward_ret")).sum())
            .group_by("date")
            .agg(long=(pl.col("long") + 1).product() - 1)
            .sort("date")
        )
        res["pos"] = pos_df.filter(pl.any_horizontal(cs.ends_with("weight") > 0))
        res["ret"] = ret_df

    return res


class BacktestEngine:

    def __init__(self,
                 data: pl.DataFrame,
                 period: str = '1d',
                 buy_fee: float = 5e-4,
                 sell_fee: float = 10e-4,):
        self.index_names: Sequence[str] = ("date", "time", "asset")
        self._cols = [*self.index_names, "price", "close", "prev_close", "open", "forward_ret"]
        self._buy_fee = buy_fee
        self._sell_fee = sell_fee
        self.data: pl.DataFrame = data.sort(self.index_names)

        self.dateList = self.data["date"].unique(maintain_order=True).to_list()
        self.timeList = self.data["time"].unique(maintain_order=True).to_list()
        self.date_num = len(self.dateList)
        self.time_num = len(self.timeList)
        self.beg_date = self.dateList[0]
        self.beg_time = self.timeList[0]

        self.limit: bool = False
        self._check()
        self._preprocess()

        # self.pos_list = list()
        self._prev_pos = defaultdict(pl.DataFrame)
        self.period = period
        self._hold_period = defaultdict(int)
        self.days = xcals.Timedelta(period).days
        if self.days < 1:
            # 小于1天的取整为1天
            logger.warning(f"{self.period} is too short, set to `1d`.")
            self.days = 1
        self._target_pos_list: list[pl.DataFrame] = list()
        self._pool: ygo.pool = None
        self.pos_list: list[pl.DataFrame] = list()
        self.pos: pl.DataFrame | None = None
        self.ret: pl.DataFrame | None = None

    def _check(self):
        # 检查数据是否合法
        check_fields = {"target_weight", "price", "close", "prev_close", "open"}
        cols = set(self.data.columns)
        assert len(check_fields & cols) >= len(check_fields), f"Missing fields({check_fields - cols}) in data."
        self.limit = "limit" in cols
        if self.limit:
            self._cols.append("limit")

    def _preprocess(self):
        # 预处理
        self.data = self.data.cast({pl.Decimal: pl.Float64}).sort(self.index_names)
        price_df_cols = [*self.index_names, "price"]
        grouper = ["date", "asset"]
        price_time_mapping = dict(prev_close="0", open="00", close="25")
        price_dfs = [
            self.data.select(price_df_cols),
            *[
                self.data.group_by(grouper).agg(time=pl.lit(t), price=pl.col(p).max()).select(price_df_cols) for p, t in
                price_time_mapping.items()
            ]
        ]

        price_df = pl.concat(price_dfs).sort(self.index_names)
        price_df = (price_df.with_columns(price_diff=pl.col("price")
                                          .diff(-1)
                                          .over(partition_by=["date", "asset"],
                                                order_by="time"))
                    .with_columns(forward_ret=pl.when(pl.col("price") > 0)
                                  .then(-pl.col("price_diff") / pl.col("price"))
                                  .otherwise(0.0))
                    .drop("price_diff"))
        self.data = price_df.join(self.data.drop("price"), on=self.index_names, how="left")
        # 归一
        self.data = self.data.with_columns(
            target_weight=pl.col("target_weight") / pl.col("target_weight").sum().over("date", "time")).filter(
            pl.col("time") < "25")
        self.data = self.data.sort(self.index_names)

    @property
    def target_weight(self) -> pl.DataFrame:
        return self.data.select(*self.index_names, "target_weight")

    def _rebalance(self, pos: pl.DataFrame, pos_id: int, date_id: int):
        """换仓"""
        date = pos["date"][0]
        time = pos["time"][0]
        # if not self.pos_list:
        cond_sub_1: bool = pos_id > 0 and pos_id == date_id  # 子仓位的第一次换仓条件
        cond_total_1: bool = pos_id == 0 and date_id == 1  # 汇总仓位的第一次换仓条件
        cond_sub_n: bool = 0 < pos_id < date_id  # 子仓位的非第一次换仓
        cond_total_n: bool = pos_id == 0 and date_id > 1  # 汇总仓位的非第一次换仓

        pos = pos.with_columns(
            target_weight=pl.when(pl.col("target_weight") < 1e-4).then(0.0).otherwise(pl.col("target_weight")))

        if date_id < pos_id and pos_id > 0:
            # 还没有到第n个子策略的开始换仓日期
            return
        elif cond_sub_1 or cond_total_1:
            # 第一次建仓
            if time <= "00":
                return
            pos = pos.with_columns(avail_sell=pl.lit(0.0), beg_weight=pl.lit(0.0))
            self._hold_period[pos_id] = 0
        elif cond_sub_n or cond_total_n:
            prev_pos = self._prev_pos[pos_id]
            prev_pos = prev_pos.select("asset", "avail_sell", beg_weight=pl.col("end_weight"))
            pos = (pos
                   .join(prev_pos, on="asset", how="full", coalesce=True)
                   .with_columns(date=pl.col("date").fill_null(date),
                                 time=pl.col("time").fill_null(time),
                                 avail_sell=pl.col("avail_sell").fill_null(0.0),
                                 beg_weight=pl.col("beg_weight").fill_null(0.0))
                   # .with_columns(target_weight=pl.col("target_weight").fill_null(pl.col("beg_weight"))) # target_weight 为 null 时，认为是不调仓
                   )
        if time <= "00":
            skip = True
        else:
            skip = 0 < self._hold_period[pos_id] < self.days if pos_id > 0 else False
        if skip:
            pos = (pos
                   .with_columns(target_weight=pl.col("beg_weight"),
                                 chg_weight=pl.lit(0.0),
                                 change_weight=pl.lit(0.0),
                                 real_weight=pl.col("beg_weight"), )
                   )
        else:
            self._hold_period[pos_id] = 0
            pos = pos.with_columns(target_weight=pl.col("target_weight").fill_null(pl.col("beg_weight")))
            target_weight = (pl.when(date=self.beg_date)
                             .then(pl.col("target_weight") / self.time_num + pl.col("beg_weight"))
                             .otherwise(pl.col("target_weight")))
            avail_sell = (pl.when(date=self.beg_date)
                          .then(0.0)
                          .when(time=self.beg_time)
                          .then(pl.col("beg_weight"))
                          .otherwise(pl.col("avail_sell"))
                          )

            pos = pos.with_columns(target_weight=target_weight, avail_sell=avail_sell)

            chg_weight = (pl.col("target_weight") - pl.col("beg_weight")).clip(lower_bound=-pl.col("avail_sell"))
            pos = pos.with_columns(chg_weight=chg_weight)
            # 涨跌停限制
            if self.limit:
                chg_weight = pl.when(pl.col("limit") > 0).then(0.0).otherwise(pl.col("chg_weight"))
                pos = pos.with_columns(chg_weight=chg_weight)
            # 更新可卖
            avail_sell = pl.col("avail_sell") + pl.col("chg_weight").clip(upper_bound=0.0)
            # 处理买入： sum(买入) <= sum(卖出)
            total_sell = ((pl.col("chg_weight") < 0) * pl.col("chg_weight")).sum()
            total_buy = ((pl.col("chg_weight") > 0) * pl.col("chg_weight")).sum()
            adj_ratio = 1.0 if pos_id <= date_id else total_sell.abs() / total_buy
            adj_buy = adj_ratio * pl.col("chg_weight")
            change_weight = pl.when(pl.col("chg_weight") > 0).then(adj_buy).otherwise(pl.col("chg_weight"))

            pos = (pos
                   .with_columns(avail_sell=avail_sell)
                   .with_columns(change_weight=change_weight)
                   .with_columns(real_weight=pl.col("beg_weight") + pl.col("change_weight"))
                   )

        pos = (
            pos
            .with_columns(end_weight=pl.col("real_weight") * (1 + pl.col("forward_ret")))
            .with_columns(end_weight=pl.col("end_weight") / pl.col("end_weight").sum())
        )

        if pos_id > 0 and self.days > 1:
            self._target_pos_list.append(
                pos.select(*self.index_names, pl.col("target_weight").alias(f"target_weight_{pos_id}")))
        else:
            self.pos_list.append(pos)

        self._prev_pos[pos_id] = pos
        self._hold_period[pos_id] += 1

        return

    def rebalance(self, date_pos: pl.DataFrame, date_id: int):

        for (time,), point_pos in date_pos.group_by("time", maintain_order=True):
            self._target_pos_list = list()
            for pos_id in range(1, 1 + self.days):
                self._pool.submit(self._rebalance)(pos=point_pos, pos_id=pos_id, date_id=date_id)
                # self._rebalance(pos=point_pos, pos_id=pos_id, date_id=date_id)
            self._pool.do()
            # 汇总所有子仓位
            if self.days > 1:
                if self._target_pos_list:
                    pos = pl.concat(self._target_pos_list, how="align").select(*self.index_names,
                                                                               target_weight=pl.sum_horizontal(
                                                                                   cs.starts_with("target_weight")))
                    # pos = pos.with_columns(target_weight=pl.when(pl.col("target_weight") < 1e-6).then(0.0).otherwise(pl.col("target_weight")))
                    # print(pos)
                    pos = point_pos.select(*self._cols, ).join(pos, on=self.index_names, how="left")
                    pos = pos.with_columns(target_weight=pl.col("target_weight") / pl.col("target_weight").sum())
                    self._rebalance(pos, 0, date_id)

    def run(self, n_jobs: int = 7, ) -> pl.DataFrame:
        pbar = tqdm(total=len(self.dateList),
                    desc="Backtesting",
                    ascii="⣀⣄⣤⣦⣶⣷⣿",
                    leave=False)
        if self._pool is None:
            self._pool = ygo.pool(n_jobs=max(n_jobs, self.days), show_progress=False, backend='threading')
        for i, date in enumerate(self.dateList):
            date_id = i + 1
            pos_df = self.data.filter(date=date)
            self.rebalance(pos_df, date_id)
            pbar.update()
        pbar.close()

        if self.pos_list:
            self.pos = pl.concat(self.pos_list).sort(self.index_names).filter(
                pl.any_horizontal(cs.ends_with("weight") > 0))
            buy_cost = (self.pos
                        .filter(pl.col("change_weight") > 0)
                        .group_by("date", "time")
                        .agg(buy_cost=(pl.col("change_weight") * self._buy_fee).sum()))
            sell_cost = (self.pos
                         .filter(pl.col("change_weight") < 0)
                         .group_by("date", "time")
                         .agg(sell_cost=(pl.col("change_weight").abs() * self._sell_fee).sum()))
            cost = (pl.concat([buy_cost, sell_cost], how="align")
                    .select("date", "time",  cost=pl.sum_horizontal(cs.ends_with("cost")).fill_null(0.0))
                    )
            self.ret = (
                self.pos
                .group_by("date", "time")
                .agg(long=(pl.col("real_weight") * pl.col("forward_ret")).sum())
                .join(cost, on=["date", "time"], how="left")
                .with_columns(long=pl.sum_horizontal(pl.col("long"), -pl.col("cost")))
                .group_by("date")
                .agg(long=(pl.col("long") + 1).product() - 1)
                .sort("date")
            )


@as_transformer
def bt(weight_df: pl.DataFrame,
       period: str = "1d",
       buy_fee: float = 5e-4,
       sell_fee: float = 10e-4,
       n_jobs: int = 7) -> BacktestEngine:
    be = BacktestEngine(data=weight_df,
                        period=period,
                        buy_fee=buy_fee,
                        sell_fee=sell_fee,)
    be.run(n_jobs=n_jobs)
    return be
