from datetime import date, datetime

import numpy as np
import polars as pl

from ldnode.calculator.schemas import EPSILON, EvaluateData, GlanceEvaluateData

PERIODS = 242


def evaluate(
    frame: pl.LazyFrame | pl.DataFrame,
    columns_dict: dict[str, str] = {
        "equity": "equity",
        "time_return": "time_return",
        "drawdown": "drawdown",
        "turnover": "turnover",
    },
) -> EvaluateData:
    """评价回测结果"""
    equity_field = columns_dict.get("equity", "equity")
    time_return_field = columns_dict.get("time_return", "time_return")
    drawdown_field = columns_dict.get("drawdown", "drawdown")
    turnover_field = columns_dict.get("turnover", None)
    res = frame.lazy() if isinstance(frame, pl.DataFrame) else frame
    exprs = [
        pl.col(equity_field).last().alias("cum_asset"),  # 累计资产
        pl.col(equity_field).first().alias("first_equity"),  # 首日资产
        pl.col(time_return_field).first().alias("first_return"),  # 首日回报
        pl.len().alias("trading_days"),  # 交易天数
        pl.col(time_return_field)
        .mean()
        .replace(0, EPSILON)
        .alias("daily_pnl_mean"),  # 日盈亏均值
        pl.col(time_return_field)
        .std(ddof=1)
        .replace(0, EPSILON)
        .fill_null(EPSILON)
        .alias("daily_pnl_std"),  # 日盈亏标准差
        pl.col(time_return_field)
        .max()
        .replace(0, EPSILON)
        .alias("daily_max_profit"),  # 日最大盈利
        pl.col(time_return_field)
        .min()
        .replace(0, EPSILON)
        .alias("daily_max_loss"),  # 日最大亏损
        pl.col(drawdown_field)
        .min()
        .replace(0, EPSILON)
        .alias("daily_max_drawdown"),  # 日最大回撤
        (pl.col(time_return_field) > 0)
        .sum()
        .fill_null(EPSILON)
        .alias("profit_count"),  # 盈利次数
        pl.col(time_return_field)
        .filter(pl.col(time_return_field) > 0)
        .mean()
        .fill_null(EPSILON)
        .alias("daily_profit_mean"),  # 日盈利均值
        pl.col(time_return_field)
        .filter(pl.col(time_return_field) < 0)
        .mean()
        .fill_null(EPSILON)
        .alias("daily_loss_mean"),  # 日亏损均值
        pl.col(time_return_field)
        .filter(pl.col(time_return_field) < 0)
        .std(ddof=1)
        .fill_null(EPSILON)
        .alias("daily_loss_std"),  # 日亏损标准差
    ]
    if turnover_field:
        exprs.append(
            pl.col(turnover_field)
            .sum()
            .truediv(pl.len())
            .fill_null(EPSILON)
            .alias("daily_turnover"),  # 日均换手率
        )
    # 计算基础统计量
    stats = res.select(exprs).collect()
    # 计算衍生指标
    capital: float = round(
        stats["first_equity"].item() / (stats["first_return"].item() + 1), 2
    )  # 初始资金
    total_return: float = stats["cum_asset"].item() / capital - 1  # 总收益率
    annu_return: float = (total_return + 1) ** (
        PERIODS / stats["trading_days"].item()
    ) - 1  # 年化收益率
    daily_win_rate: float = (
        stats["profit_count"].item() / stats["trading_days"].item()
    )  # 日胜率

    daily_profit_loss_rate: float = (
        -stats["daily_profit_mean"].item() / stats["daily_loss_mean"].item()
    )  # 日盈亏比
    sharpe_ratio: float = (
        stats["daily_pnl_mean"].item() / stats["daily_pnl_std"].item()
    ) * np.sqrt(PERIODS)  # 夏普比率
    sortino_ratio: float = (
        stats["daily_pnl_mean"].item() / stats["daily_loss_std"].item()
    ) * np.sqrt(PERIODS)  # 索提诺比率
    calmar_ratio: float = (
        annu_return / -stats["daily_max_drawdown"].item()
    )  # 卡玛比比率

    daily_max_drawdown = stats["daily_max_drawdown"].item()  # 最大回撤
    daily_max_drawdown_end_index = (
        res.select(pl.col(drawdown_field).arg_min()).collect().item()
    )  # 最大回撤结束日期索引(如果有多处最大回撤，选最先出现的索引)
    _zero_drawdown_index_df = res.select(
        pl.arg_where(pl.col(drawdown_field) == 0).alias("zero_drawdown_index")
    ).collect()  # 最大回撤为0的日期索引df
    daily_max_drawdown_start_index = _zero_drawdown_index_df.select(
        pl.col("zero_drawdown_index")
        .filter(pl.col("zero_drawdown_index") < daily_max_drawdown_end_index)
        .last()
        .fill_null(0)
    ).item()  # 最大回撤开始日期索引
    daily_max_drawdown_duration_trade_days: int = (
        daily_max_drawdown_end_index - daily_max_drawdown_start_index
    )  # 最大回撤持续交易日
    daily_max_drawdown_start_date_lf = res.select(
        pl.col("trade_date").gather(daily_max_drawdown_start_index)
    )  # 最大回撤开始日期
    daily_max_drawdown_end_date_lf = res.select(
        pl.col("trade_date").gather(daily_max_drawdown_end_index)
    )  # 最大回撤结束日期
    daily_max_drawdown_start_date_df, daily_max_drawdown_end_date_df = pl.collect_all(
        [daily_max_drawdown_start_date_lf, daily_max_drawdown_end_date_lf]
    )
    daily_max_drawdown_start_date = (
        daily_max_drawdown_start_date_df.item()
    )  # 最大回撤开始日期
    daily_max_drawdown_end_date = (
        daily_max_drawdown_end_date_df.item()
    )  # 最大回撤结束日期
    daily_max_drawdown_duration_calendar_days = (
        daily_max_drawdown_end_date - daily_max_drawdown_start_date
    ).days  # 最大回撤持续日历天数
    evaluate_dict: EvaluateData = {
        "cum_asset": stats["cum_asset"].item(),
        "total_return": total_return,
        "trading_days": stats["trading_days"].item(),
        "annu_return": annu_return,
        "daily_turnover": stats["daily_turnover"].item() if turnover_field else None,
        "profit_count": stats["profit_count"].item(),
        "daily_win_rate": daily_win_rate,
        "daily_profit_mean": stats["daily_profit_mean"].item(),
        "daily_loss_mean": stats["daily_loss_mean"].item(),
        "daily_profit_loss_rate": daily_profit_loss_rate,
        "daily_pnl_mean": stats["daily_pnl_mean"].item(),
        "daily_pnl_std": stats["daily_pnl_std"].item(),
        "daily_loss_std": stats["daily_loss_std"].item(),
        "daily_max_profit": stats["daily_max_profit"].item(),
        "daily_max_loss": stats["daily_max_loss"].item(),
        "daily_max_drawdown": daily_max_drawdown,
        "daily_max_drawdown_duration_trade_days": daily_max_drawdown_duration_trade_days,
        "daily_max_drawdown_duration_calendar_days": daily_max_drawdown_duration_calendar_days,
        "sharpe_ratio": sharpe_ratio,
        "calmar_ratio": calmar_ratio,
        "sortino_ratio": sortino_ratio,
    }
    return {
        k: round(v, 4) if isinstance(v, float) else v for k, v in evaluate_dict.items()
    }  # type: ignore # 保留4位小数


def glance_evaluate(
    frame: pl.LazyFrame | pl.DataFrame,
    start_date: str,
    end_date: str,
    created_date: str,
    columns_dict: dict[str, str] = {
        "equity": "equity",
        "time_return": "time_return",
        "drawdown": "drawdown",
    },
) -> GlanceEvaluateData:
    """概览评价回测结果"""
    res = frame.lazy() if isinstance(frame, pl.DataFrame) else frame
    equity_field = columns_dict.get("equity", "equity")
    time_return_field = columns_dict.get("time_return", "time_return")
    drawdown_field = columns_dict.get("drawdown", "drawdown")
    this_year_start_date: date = datetime.now().date().replace(month=1, day=1)
    this_month_start_date: date = datetime.now().date().replace(day=1)

    exprs = [
        pl.col(equity_field).last().alias("cum_asset"),  # 累计资产
        pl.col(equity_field).first().alias("first_equity"),  # 首日资产
        pl.col(time_return_field).first().alias("first_return"),  # 首日回报
        pl.len().alias("trading_days"),  # 交易天数
        pl.col(time_return_field)
        .mean()
        .replace(0, EPSILON)
        .alias("daily_pnl_mean"),  # 日盈亏均值
        pl.col(time_return_field)
        .std(ddof=1)
        .replace(0, EPSILON)
        .fill_null(EPSILON)
        .alias("daily_pnl_std"),  # 日盈亏标准差
        pl.col(drawdown_field)
        .min()
        .replace(0, EPSILON)
        .alias("daily_max_drawdown"),  # 日最大回撤
    ]
    # 计算基础统计量
    (
        stats,
        last_1_return_df,
        last_5_return_df,
        last_20_return_df,
        last_60_return_df,
        last_242_return_df,
        this_month_return_df,
        this_year_return_df,
        created_return_df,
    ) = pl.collect_all(
        [
            res.select(exprs),
            res.select(pl.col(time_return_field).get(-1)),
            res.select(
                pl.col(equity_field).get(-1) / pl.col(equity_field).shift(5).get(-1) - 1
            ),
            res.select(
                pl.col(equity_field).get(-1) / pl.col(equity_field).shift(20).get(-1)
                - 1
            ),
            res.select(
                pl.col(equity_field).get(-1) / pl.col(equity_field).shift(60).get(-1)
                - 1
            ),
            res.select(
                pl.col(equity_field).get(-1) / pl.col(equity_field).shift(242).get(-1)
                - 1
            ),
            res.select(
                (
                    pl.col(equity_field).last()
                    / pl.coalesce(
                        pl.col(equity_field)
                        .filter(pl.col("trade_date") < this_month_start_date)
                        .last(),
                        pl.col(equity_field).first(),
                    )
                )
                - 1
            ),
            res.select(
                (
                    pl.col(equity_field).last()
                    / pl.coalesce(
                        pl.col(equity_field)
                        .filter(pl.col("trade_date") < this_year_start_date)
                        .last(),
                        pl.col(equity_field).first(),
                    )
                )
                - 1
            ),
            res.select(
                pl.col(equity_field).get(-1)
                / pl.col(equity_field)
                .filter(
                    pl.col("trade_date")
                    >= datetime.strptime(created_date, "%Y-%m-%d").date()
                )
                .first()
                - 1
            ),
        ]
    )
    # 安全 round：值为 None 时保留 None（表示该指标无法计算），前端显示为 "-"
    def _round(v, n=4):
        return round(v, n) if v is not None else None

    # 计算衍生指标
    _first_equity = stats["first_equity"].item()
    _first_return = stats["first_return"].item()
    _cum_asset = stats["cum_asset"].item()
    _trading_days = stats["trading_days"].item()
    if _first_equity is None or _first_return is None or _cum_asset is None or not _trading_days:
        raise ValueError("回测结果为空，无法计算收益指标（策略可能在所有交易日均无有效选股）")
    capital: float = round(_first_equity / (_first_return + 1), 2)
    total_return: float = _cum_asset / capital - 1
    annu_return: float = (total_return + 1) ** (PERIODS / _trading_days) - 1
    _daily_max_drawdown = stats["daily_max_drawdown"].item()
    _pnl_mean = stats["daily_pnl_mean"].item()
    _pnl_std = stats["daily_pnl_std"].item()
    sharpe_ratio = (_pnl_mean / _pnl_std) * np.sqrt(PERIODS) if _pnl_mean is not None and _pnl_std else None

    evaluate_dict: GlanceEvaluateData = {
        "start_date": start_date,
        "end_date": end_date,
        "created_date": created_date,
        "annu_return": _round(annu_return),
        "daily_max_drawdown": _round(_daily_max_drawdown),
        "sharpe_ratio": _round(sharpe_ratio),
        "last_1_return": _round(last_1_return_df.item()),
        "last_5_return": _round(last_5_return_df.item()),
        "last_20_return": _round(last_20_return_df.item()),
        "last_60_return": _round(last_60_return_df.item()),
        "last_242_return": _round(last_242_return_df.item()),
        "this_month_return": _round(this_month_return_df.item()),
        "this_year_return": _round(this_year_return_df.item()),
        "created_return": _round(created_return_df.item()),
        "total_return": _round(total_return),
    }
    return evaluate_dict
