import polars as pl


def gen_total_return_expr(equity_expr: pl.Expr, capital: float) -> pl.Expr:
    """总收益率表达式"""
    return equity_expr / capital - 1


def gen_drawdown_expr(equity_expr: pl.Expr) -> pl.Expr:
    """回撤表达式"""
    return (equity_expr - equity_expr.cum_max()) / equity_expr.cum_max()


def gen_equity_expr(time_return: pl.Expr, capital: float) -> pl.Expr:
    """资产表达式"""
    return (time_return.fill_null(0) + 1).cum_prod() * capital


def planning_earnings_curve(
    res_lf: pl.LazyFrame,
    capital: float,
    with_benchmark: bool = True,
) -> pl.LazyFrame:
    """规划生成收益曲线"""
    # 总收益率
    total_return_expr = gen_total_return_expr(pl.col("equity"), capital)
    # 策略回撤
    drawdown_expr = gen_drawdown_expr(pl.col("equity"))
    # 基准净值
    equity_benchmark_expr = gen_equity_expr(pl.col("time_return_benchmark"), capital)
    # 基准总收益率
    total_return_benchmark_expr = gen_total_return_expr(equity_benchmark_expr, capital)
    # 基准回撤
    drawdown_benchmark_expr = gen_drawdown_expr(equity_benchmark_expr)
    # 相对超额收益率
    time_return_relative_excess_expr = (pl.col("time_return") + 1) / (
        pl.col("time_return_benchmark") + 1
    ) - 1
    # 相对超额净值
    equity_relative_excess_expr = gen_equity_expr(
        time_return_relative_excess_expr, capital
    )
    # 相对超额总收益率
    total_return_relative_excess_expr = gen_total_return_expr(
        equity_relative_excess_expr, capital
    )
    # 相对超额回撤
    drawdown_relative_excess_expr = gen_drawdown_expr(equity_relative_excess_expr)
    # 绝对超额收益率
    time_return_absolute_excess_expr = pl.col("time_return") - pl.col(
        "time_return_benchmark"
    )
    # 绝对超额净值
    equity_absolute_excess_expr = gen_equity_expr(
        time_return_absolute_excess_expr, capital
    )
    # 绝对超额总收益率
    total_return_absolute_excess_expr = gen_total_return_expr(
        equity_absolute_excess_expr, capital
    )
    # 绝对超额回撤
    drawdown_absolute_excess_expr = gen_drawdown_expr(equity_absolute_excess_expr)

    current_gen_exprs = [
        total_return_expr.alias("total_return"),
        drawdown_expr.alias("drawdown"),
    ]
    benchmark_gen_exprs = [
        equity_benchmark_expr.alias("equity_benchmark"),
        total_return_benchmark_expr.alias("total_return_benchmark"),
        drawdown_benchmark_expr.alias("drawdown_benchmark"),
        time_return_relative_excess_expr.alias("time_return_relative_excess"),
        time_return_absolute_excess_expr.alias("time_return_absolute_excess"),
        equity_relative_excess_expr.alias("equity_relative_excess"),
        equity_absolute_excess_expr.alias("equity_absolute_excess"),
        total_return_relative_excess_expr.alias("total_return_relative_excess"),
        total_return_absolute_excess_expr.alias("total_return_absolute_excess"),
        drawdown_relative_excess_expr.alias("drawdown_relative_excess"),
        drawdown_absolute_excess_expr.alias("drawdown_absolute_excess"),
    ]
    gen_exprs = (
        [
            *current_gen_exprs,
            *benchmark_gen_exprs,
        ]
        if with_benchmark
        else current_gen_exprs
    )
    current_select_exprs = [
        pl.col("trade_date"),
        pl.col("equity").round(2),
        pl.col("time_return").round(4),
        pl.col("total_return").round(4),
        pl.col("drawdown").round(4),
        pl.col("turnover").round(4),
        pl.col("commission").round(2),
        pl.col("hold_securities"),
    ]
    benchmark_select_exprs = [
        pl.col("equity_benchmark").round(2),
        pl.col("time_return_benchmark").round(4),
        pl.col("total_return_benchmark").round(4),
        pl.col("drawdown_benchmark").round(4),
        pl.col("time_return_relative_excess").round(4),
        pl.col("equity_relative_excess").round(2),
        pl.col("total_return_relative_excess").round(4),
        pl.col("drawdown_relative_excess").round(4),
        pl.col("time_return_absolute_excess").round(4),
        pl.col("equity_absolute_excess").round(2),
        pl.col("total_return_absolute_excess").round(4),
        pl.col("drawdown_absolute_excess").round(4),
    ]
    select_exprs = (
        [
            *current_select_exprs,
            *benchmark_select_exprs,
        ]
        if with_benchmark
        else current_select_exprs
    )
    return res_lf.with_columns(gen_exprs).select(select_exprs)
