# Copyright (c) ZhangYundi.
# Licensed under the MIT License. 
# Created on 2025/12/23 15:41
# Description:
import math
from typing import Sequence, Literal

import polars as pl
from ..transformer import as_transformer
import numpy as np
import logair
import cvxpy as cp
import polars.selectors as cs
from tqdm.auto import tqdm

@as_transformer
def qcut(score_df: pl.DataFrame, N: int = 10, grouper: Sequence[str] | None = None) -> pl.DataFrame:
    grouper = ["date", "time"] if not grouper else ["date", "time", *grouper]
    labels = [str(i) for i in range(1, N + 1)]
    cat = pl.Enum(labels)
    return (
        score_df
        .drop_nulls(grouper)
        .with_columns(
            pl.col("score")
            .qcut(N,
                  labels=labels,
                  allow_duplicates=True)
            .over(grouper)
            .cast(cat)
            .alias("quantile"))
        .cast({pl.Categorical: pl.Utf8})
        .filter(quantile=str(N))
        .drop("quantile")
        .with_columns(target_weight=(1/pl.col("asset").n_unique()).over("date", "time"))
    )

@as_transformer
def top_k(score_df: pl.DataFrame, num: int = 500, grouper: Sequence[str] | None = None) -> pl.DataFrame:
    """选择 top num 数量"""
    grouper = ["date", "time"] if not grouper else ["date", "time", *grouper]
    return (
        score_df
        .drop_nulls(grouper)
        .group_by(grouper, maintain_order=True)
        .agg(pl.all().top_k_by(k=num, by="score"))
        .explode(pl.all().exclude(grouper))
        .with_columns(target_weight=(1/pl.col("asset").n_unique()).over("date", "time"))
    )

def mfe(score_df: pl.DataFrame,
        weight_bias: float = -1.0,
        indust_bias: float = -1.0,
        style_bias: dict[str, float] | None = None,
        weight_limit: float = -1.0,
        to_limit: float = -1.0,
        components_ratio: float = -1.0,
        index: Sequence[str] = ("date", "time", "asset"),
        check_nan_cols: Sequence[str] | None = None,
        bm_w_name: str = "bm_w") -> pl.DataFrame:
    """
    最大因子暴露组合

    Parameters
    ----------
    score_df: polars.DataFrame
    weight_bias: float
        最大权重偏离, 偏离绝对值，| weight - bm_w | <= max_weight_bias, 小于0则不做约束
    weight_limit: float
        个股最大权重约束, 小于0则不做约束
    to_limit: float
        组合换手约束
    components_ratio: float
        成分股权重约束
    indust_bias: float
        行业偏离百分比
    style_bias: dict[str, float] | None
        风格约束
    index: Sequence[str]
    check_nan_cols: Sequence[str] | None
        检查空值的列
    bm_w_name: str
        基准个股权重列名

    Returns
    -------
    polars.DataFrame

    """

    logger = logair.get_logger(__name__)

    bm_w = pl.col(bm_w_name)
    score_df = (score_df
                .with_columns(bm_w.fill_null(0.0).fill_nan(0.0).alias(bm_w_name))
                .with_columns((bm_w / bm_w.sum()).alias(bm_w_name)))
    score_df = score_df.select(cs.by_name(*check_nan_cols, require_all=False)).drop_nulls().drop_nans()
    raw_index = score_df.select(index)
    stock_num = score_df["asset"].n_unique()

    bm_w_arr: np.ndarray = score_df[bm_w_name].to_numpy()

    #### 构建优化问题
    w_dev = cp.Variable(stock_num)
    bm_w_arr = bm_w_arr.reshape(w_dev.shape)
    # 初始解
    initial_guess = np.full(shape=w_dev.shape, fill_value=0)
    w_dev.value = initial_guess

    w = w_dev + bm_w_arr

    #### 约束1: 非空+现金中性
    constraints = [w >= 0, cp.sum(w_dev) == 0, ]
    #### 约束2: 最大权重偏离
    if weight_bias >= 0:
        constraints.append(cp.abs(w_dev) <= weight_bias)
    ### 约束3: 最大权重约束
    if weight_limit >= 0:
        constraints.append(w <= weight_limit)
    ### 约束4: 换手约束
    if to_limit >= 0 and "prev_weight" in score_df.columns:
        prev_w_arr: np.ndarray = score_df["prev_weight"].to_numpy()
        constraints.append(cp.sum(cp.abs(w-prev_w_arr)) <= to_limit)
    # constraints.append(cp.sum(bm_w.values.reshape(w_dev.shape) + w_dev) <= 1,)
    #### 约束5: 最小成分股占比
    if components_ratio > 0:
        bm_codes = score_df.select(is_components=pl.when(bm_w > 0).then(1).otherwise(0)).to_numpy().reshape(w_dev.shape)
        constraints.append(bm_codes.T @ w >= components_ratio)
    #### 约束6: 行业偏离
    if indust_bias >= 0:
        industry_exposure = score_df.select("group").to_dummies().to_numpy()
        indust_bm = industry_exposure.T @ bm_w_arr.reshape(w_dev.shape)
        constraints.append(cp.abs(industry_exposure.T @ w_dev) <= indust_bias * indust_bm)
    #### 约束7: 风格偏离
    #######################################################################
    #                               风格约束                                #
    #######################################################################
    if style_bias is not None:
        bm_score = score_df.filter(bm_w > 0)
        for style_name, bias in style_bias.items():
            if bias < 0:
                continue
            style_mu = (bm_score[style_name] * bm_score[bm_w_name]).sum()
            style_variance = (bm_score[bm_w_name] * (bm_score[style_name] - style_mu) ** 2).sum()
            style_sigma = math.sqrt(style_variance)
            style_norm = (score_df[style_name] - style_mu) / style_sigma
            style_norm = style_norm.to_numpy().reshape(w_dev.shape)
            constraints.append(
                cp.abs(style_norm.T @ w) <= bias
            )

    #### 求解目标
    score = score_df["score"].to_numpy().reshape(w_dev.shape)
    alpha = score.T @ w
    objective = cp.Maximize(alpha)
    problem = cp.Problem(objective, constraints)
    try:
        problem.solve(**dict(
            solver=cp.ECOS,
            # solver=cp.CLARABEL,
            feastol=1e-4,  # 增大可行性容忍度
            abstol=1e-4,  # 增大绝对容忍度
            reltol=1e-4,  # 增大相对容忍度
            verbose=False,
        ))
    except Exception as e:
        logger.error(e)
        return
    w_dev = w_dev.value
    if w_dev is None:
        logger.error(f"Optimize failed: {score_df["date"][0]}\n{score_df.drop(*index, "group", strict=False).describe()}",)
        return
    w_dev = np.where(np.isinf(w_dev), 0.0, w_dev)
    w = score_df.select(*index, bm_w_name, w_dev=w_dev)
    w_arr = (w[bm_w_name] + w["w_dev"]).fill_null(0.0).fill_nan(0.0).to_numpy()
    w_arr = np.where(w_arr < 1e-6, 0.0, w_arr)
    w_arr = w_arr / w_arr.sum()
    return raw_index.with_columns(target_weight=w_arr).filter(pl.col("target_weight") > 0.0)


@as_transformer
def MFE(score_df: pl.DataFrame,
        weight_bias: float = -1.0,
        indust_bias: float = -1.0,
        style_bias: dict[str, float] | None | float = -1.0,
        weight_limit: float = -1.0,
        to_limit: float = -1.0,
        components_ratio: float = -1.0,
        bm_w_name: str = "bm_w",) -> pl.DataFrame:
    """
    最大因子暴露组合:
        - bm_w: 基准个股权重
        - 风格因子: 以 Barra 结尾
        - group: 所属行业

    Parameters
    ----------
    score_df: polars.DataFrame
    weight_bias: float
        最大权重偏离, 偏离绝对值，| weight - bm_w | <= max_weight_bias, 小于0则不做约束
    weight_limit: float
        个股最大权重约束, 小于0则不做约束, 默认不约束
    to_limit: float
        换手约束, 小于0则不做约束，默认不约束
    components_ratio: float
        成分股权重约束, 小于0则不做约束,默认不约束
    indust_bias: float
        行业偏离百分比, 小于0则不做约束, 默认不约束
    style_bias: dict[str, float] | None | float
        风格约束: float 表示所有风格做同样的约束，默认不约束
    bm_w_name: str
        基准个股权重列名

    Returns
    -------
    polars.DataFrame

    """
    index = {"date", "time", "asset"} & set(score_df.columns)
    index = list(index)
    score_df = score_df.drop_nulls([*index, "score", "group"]).drop_nans("score")
    check_nan_cols = [*index, "score", bm_w_name]
    bm_w = pl.col(bm_w_name)
    score_df = score_df.with_columns((bm_w.fill_null(0.0).fill_nan(0.0)).alias(bm_w_name))

    if isinstance(style_bias, (float, int)):
        if style_bias < 0:
            style_bias = None
        else:
            style_names = score_df.select(cs.ends_with("Barra")).columns
            style_bias = {style_name: style_bias for style_name in style_names}
    if style_bias is not None:
        check_nan_cols.extend(style_bias.keys())
    if indust_bias >= 0:
        check_nan_cols.append("group")
    if to_limit >= 0:
        check_nan_cols.append("prev_weight")
    results = list()
    prev_w: pl.DataFrame | None = None
    pbar = tqdm(total=score_df["date"].n_unique(), desc="MFE", ascii="⣀⣄⣤⣦⣶⣷⣿",)
    for (date, ), score_df_ in score_df.group_by("date", maintain_order=True):
        pbar.set_postfix_str(date)
        if prev_w is not None:
            score_df_ = score_df_.join(prev_w.drop("date").rename({"target_weight": "prev_weight"}),
                                       on="asset",
                                       how="left").with_columns(pl.col("prev_weight").fill_null(0.0))
        w_df = mfe(score_df=score_df_,
                   weight_bias=weight_bias,
                   weight_limit=weight_limit,
                   to_limit=to_limit,
                   indust_bias=indust_bias,
                   components_ratio=components_ratio,
                   style_bias=style_bias,
                   index=index,
                   check_nan_cols=check_nan_cols,
                   bm_w_name=bm_w_name, )
        if w_df is None:
            w_df = prev_w
        w_df: pl.DataFrame = w_df.with_columns(date=pl.lit(date)).drop_nulls()
        results.append(w_df)
        prev_w = w_df
        pbar.update()
    return pl.concat(results)