"""Point-in-time cross-sectional factor research tools.

This module deliberately accepts already assembled panels instead of reading the
latest processed files.  A caller must attach the date on which every factor
was available.  That makes the availability check explicit and prevents a
current financial snapshot from silently leaking into historical backtests.
All returns are decimal returns (``0.01`` means one percent).
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Sequence

import numpy as np
import pandas as pd


@dataclass(frozen=True)
class FactorBacktestResult:
    """Outputs from an independent single-factor cross-sectional backtest."""

    normalized_panel: pd.DataFrame
    group_returns: pd.DataFrame
    period_returns: pd.DataFrame
    ic: pd.DataFrame
    summary: pd.DataFrame


@dataclass(frozen=True)
class WeightOptimizationResult:
    """Walk-forward weights and the strictly out-of-sample composite test."""

    weights: pd.Series
    weight_history: pd.DataFrame
    oos_panel: pd.DataFrame
    oos_backtest: FactorBacktestResult


_GROUP_COLUMNS = ["rebalance_date", "factor", "group", "group_return", "member_count"]
_PERIOD_COLUMNS = [
    "rebalance_date", "factor", "long_return", "short_return", "long_short_gross_return",
    "long_turnover", "short_turnover", "long_short_turnover", "long_only_net_return",
    "long_short_net_return", "long_count", "short_count",
]
_IC_COLUMNS = ["rebalance_date", "factor", "ic", "rank_ic", "member_count"]
_SUMMARY_COLUMNS = [
    "factor", "periods", "mean_ic", "mean_rank_ic", "ic_positive_rate",
    "mean_long_short_gross_return", "mean_long_short_net_return",
    "cumulative_long_short_net_return", "mean_long_only_net_return",
]


def _require_columns(frame: pd.DataFrame, columns: Sequence[str]) -> None:
    missing = [column for column in columns if column not in frame.columns]
    if missing:
        raise ValueError(f"因子面板缺少必要字段：{', '.join(missing)}")


def _validate_options(quantiles: int, winsor_limits: tuple[float, float], min_assets: int) -> None:
    if quantiles < 2:
        raise ValueError("quantiles 必须至少为 2")
    if min_assets < quantiles:
        raise ValueError("min_assets 不能小于 quantiles")
    lower, upper = winsor_limits
    if not (0.0 <= lower < upper <= 1.0):
        raise ValueError("winsor_limits 必须满足 0 <= lower < upper <= 1")


def prepare_factor_panel(
    panel: pd.DataFrame,
    factor_columns: Sequence[str],
    *,
    rebalance_col: str = "rebalance_date",
    security_col: str = "source_code",
    available_date_col: str = "available_date",
    forward_return_col: str = "forward_return",
    winsor_limits: tuple[float, float] = (0.01, 0.99),
) -> pd.DataFrame:
    """Validate PIT provenance and cross-sectionally winsorize and z-score factors.

    ``available_date`` must be the first tradable date for the information, not
    a report period or a data-download date.  Values that are unavailable on a
    rebalance date raise instead of being shifted implicitly.
    """
    if not factor_columns:
        raise ValueError("至少需要一个因子字段")
    _validate_options(2, winsor_limits, 2)
    required = [rebalance_col, security_col, available_date_col, forward_return_col, *factor_columns]
    _require_columns(panel, required)
    frame = panel.copy()
    frame[rebalance_col] = pd.to_datetime(frame[rebalance_col], errors="coerce")
    frame[available_date_col] = pd.to_datetime(frame[available_date_col], errors="coerce")
    if frame[rebalance_col].isna().any() or frame[available_date_col].isna().any():
        raise ValueError("rebalance_date 与 available_date 不得为空且必须是有效日期")
    if frame[security_col].isna().any():
        raise ValueError("security 标识不得为空")
    duplicate = frame.duplicated([rebalance_col, security_col], keep=False)
    if duplicate.any():
        raise ValueError("每个调仓日每只证券只能有一行因子观测值")
    leaked = frame[available_date_col] > frame[rebalance_col]
    if leaked.any():
        examples = frame.loc[leaked, [security_col, available_date_col, rebalance_col]].head(5)
        raise ValueError(f"检测到未来函数：available_date 晚于 rebalance_date：{examples.to_dict('records')}")
    frame[forward_return_col] = pd.to_numeric(frame[forward_return_col], errors="coerce")
    for factor in factor_columns:
        frame[factor] = pd.to_numeric(frame[factor], errors="coerce")

    lower, upper = winsor_limits
    for factor in factor_columns:
        winsorized = pd.Series(np.nan, index=frame.index, dtype=float)
        zscore = pd.Series(np.nan, index=frame.index, dtype=float)
        for _, index in frame.groupby(rebalance_col, sort=False).groups.items():
            values = frame.loc[index, factor]
            valid = values[np.isfinite(values)]
            if valid.empty:
                continue
            clipped = values.clip(lower=valid.quantile(lower), upper=valid.quantile(upper))
            winsorized.loc[index] = clipped
            usable = clipped[np.isfinite(clipped)]
            std = float(usable.std(ddof=0))
            if len(usable) > 1 and np.isfinite(std) and std > 0:
                zscore.loc[usable.index] = (usable - float(usable.mean())) / std
            elif len(usable) == 1:
                zscore.loc[usable.index] = 0.0
            else:
                zscore.loc[usable.index] = 0.0
        frame[f"{factor}_winsorized"] = winsorized
        frame[f"{factor}_zscore"] = zscore
    return frame.sort_values([rebalance_col, security_col]).reset_index(drop=True)


def _quantile_groups(values: pd.Series, quantiles: int) -> pd.Series:
    """Assign stable cross-sectional buckets even where factor values tie."""
    ranked = values.rank(method="first")
    return pd.qcut(ranked, q=quantiles, labels=False, duplicates="drop").astype(int) + 1


def _turnover(current: set[Any], previous: set[Any] | None) -> float:
    if previous is None:
        return 1.0
    if not current and not previous:
        return 0.0
    current_weight = 1.0 / len(current) if current else 0.0
    previous_weight = 1.0 / len(previous) if previous else 0.0
    names = current | previous
    return 0.5 * sum(abs((current_weight if name in current else 0.0) -
                         (previous_weight if name in previous else 0.0)) for name in names)


def _pearson_ic(left: pd.Series, right: pd.Series) -> float:
    """Correlation without an optional SciPy dependency."""
    valid = pd.DataFrame({"left": left, "right": right}).dropna()
    if len(valid) < 2 or valid["left"].std(ddof=0) == 0 or valid["right"].std(ddof=0) == 0:
        return float("nan")
    return float(valid["left"].corr(valid["right"]))


def _rank_ic(left: pd.Series, right: pd.Series) -> float:
    """Spearman correlation computed as Pearson correlation of average ranks."""
    valid = pd.DataFrame({"left": left, "right": right}).dropna()
    return _pearson_ic(valid["left"].rank(method="average"), valid["right"].rank(method="average"))


def _empty_backtest(normalized_panel: pd.DataFrame) -> FactorBacktestResult:
    return FactorBacktestResult(
        normalized_panel=normalized_panel,
        group_returns=pd.DataFrame(columns=_GROUP_COLUMNS),
        period_returns=pd.DataFrame(columns=_PERIOD_COLUMNS),
        ic=pd.DataFrame(columns=_IC_COLUMNS),
        summary=pd.DataFrame(columns=_SUMMARY_COLUMNS),
    )


def run_factor_backtest(
    panel: pd.DataFrame,
    factor_columns: Sequence[str],
    *,
    rebalance_col: str = "rebalance_date",
    security_col: str = "source_code",
    available_date_col: str = "available_date",
    forward_return_col: str = "forward_return",
    quantiles: int = 5,
    transaction_cost_bps: float = 0.0,
    long_short: bool = True,
    winsor_limits: tuple[float, float] = (0.01, 0.99),
    min_assets: int = 10,
) -> FactorBacktestResult:
    """Backtest each factor with PIT-safe cross-sectional portfolios.

    The top bucket is long and the bottom bucket is short.  Transaction costs
    are charged on one-way turnover at each rebalance, including the initial
    portfolio entry.  If no date has enough valid names, empty result frames
    are returned rather than a misleading zero-return backtest.
    """
    _validate_options(quantiles, winsor_limits, min_assets)
    if transaction_cost_bps < 0:
        raise ValueError("transaction_cost_bps 不能为负数")
    normalized = prepare_factor_panel(
        panel, factor_columns, rebalance_col=rebalance_col, security_col=security_col,
        available_date_col=available_date_col, forward_return_col=forward_return_col,
        winsor_limits=winsor_limits,
    )
    group_rows: list[dict[str, Any]] = []
    period_rows: list[dict[str, Any]] = []
    ic_rows: list[dict[str, Any]] = []
    cost_rate = transaction_cost_bps / 10_000.0

    for factor in factor_columns:
        zscore_col = f"{factor}_zscore"
        previous_long: set[Any] | None = None
        previous_short: set[Any] | None = None
        for rebalance_date, slice_ in normalized.groupby(rebalance_col, sort=True):
            usable = slice_.loc[
                np.isfinite(slice_[zscore_col]) & np.isfinite(slice_[forward_return_col]),
                [security_col, zscore_col, forward_return_col],
            ].copy()
            if len(usable) < min_assets:
                continue
            usable["_group"] = _quantile_groups(usable[zscore_col], quantiles)
            grouped = usable.groupby("_group", sort=True)[forward_return_col].agg(["mean", "size"])
            for bucket, values in grouped.iterrows():
                group_rows.append({
                    "rebalance_date": rebalance_date, "factor": factor, "group": int(bucket),
                    "group_return": float(values["mean"]), "member_count": int(values["size"]),
                })
            long_members = set(usable.loc[usable["_group"].eq(quantiles), security_col])
            short_members = set(usable.loc[usable["_group"].eq(1), security_col])
            # qcut only drops bins for invalid inputs; valid n >= quantiles guarantees both ends.
            long_return = float(usable.loc[usable["_group"].eq(quantiles), forward_return_col].mean())
            short_return = float(usable.loc[usable["_group"].eq(1), forward_return_col].mean())
            long_turnover = _turnover(long_members, previous_long)
            short_turnover = _turnover(short_members, previous_short)
            long_short_turnover = 0.5 * (long_turnover + short_turnover)
            gross = long_return - short_return
            period_rows.append({
                "rebalance_date": rebalance_date,
                "factor": factor,
                "long_return": long_return,
                "short_return": short_return,
                "long_short_gross_return": gross,
                "long_turnover": long_turnover,
                "short_turnover": short_turnover,
                "long_short_turnover": long_short_turnover,
                "long_only_net_return": long_return - long_turnover * cost_rate,
                "long_short_net_return": gross - long_short_turnover * cost_rate if long_short else np.nan,
                "long_count": len(long_members),
                "short_count": len(short_members),
            })
            ic_rows.append({
                "rebalance_date": rebalance_date,
                "factor": factor,
                "ic": _pearson_ic(usable[zscore_col], usable[forward_return_col]),
                "rank_ic": _rank_ic(usable[zscore_col], usable[forward_return_col]),
                "member_count": len(usable),
            })
            previous_long, previous_short = long_members, short_members

    if not period_rows:
        return _empty_backtest(normalized)
    group_returns = pd.DataFrame(group_rows, columns=_GROUP_COLUMNS)
    period_returns = pd.DataFrame(period_rows, columns=_PERIOD_COLUMNS)
    ic = pd.DataFrame(ic_rows, columns=_IC_COLUMNS)
    summaries = []
    for factor, periods in period_returns.groupby("factor", sort=True):
        factor_ic = ic.loc[ic["factor"].eq(factor)]
        net = periods["long_short_net_return"] if long_short else periods["long_only_net_return"]
        summaries.append({
            "factor": factor,
            "periods": int(len(periods)),
            "mean_ic": float(factor_ic["ic"].mean()),
            "mean_rank_ic": float(factor_ic["rank_ic"].mean()),
            "ic_positive_rate": float(factor_ic["rank_ic"].gt(0).mean()),
            "mean_long_short_gross_return": float(periods["long_short_gross_return"].mean()),
            "mean_long_short_net_return": float(periods["long_short_net_return"].mean()),
            "cumulative_long_short_net_return": float((1.0 + net.dropna()).prod() - 1.0),
            "mean_long_only_net_return": float(periods["long_only_net_return"].mean()),
        })
    return FactorBacktestResult(
        normalized_panel=normalized,
        group_returns=group_returns,
        period_returns=period_returns,
        ic=ic,
        summary=pd.DataFrame(summaries, columns=_SUMMARY_COLUMNS),
    )


def _weight_candidates(count: int, candidate_count: int, random_state: int,
                       max_weight: float = 1.0) -> np.ndarray:
    if candidate_count < count + 1:
        raise ValueError("candidate_count 至少应为因子数加一")
    if max_weight < 1.0 / count or max_weight > 1.0:
        raise ValueError("max_weight 必须不小于等权权重且不大于 1")
    rng = np.random.default_rng(random_state)
    candidates: list[np.ndarray] = [np.full(count, 1.0 / count)]
    if max_weight >= 1.0:
        candidates.extend(np.eye(count)[index] for index in range(count))
    while len(candidates) < candidate_count:
        batch = rng.dirichlet(np.ones(count), size=max(64, candidate_count * 2))
        candidates.extend(row for row in batch if row.max() <= max_weight + 1e-12)
    return np.asarray(candidates[:candidate_count], dtype=float)


def _score_weights(
    prepared: pd.DataFrame,
    factor_columns: Sequence[str],
    weights: np.ndarray,
    dates: Sequence[pd.Timestamp],
    *,
    rebalance_col: str,
    forward_return_col: str,
    min_assets: int,
    quantiles: int,
    objective: str,
    date_slices: dict[pd.Timestamp, pd.DataFrame] | None = None,
) -> tuple[float, int]:
    zscore_columns = [f"{factor}_zscore" for factor in factor_columns]
    values: list[float] = []
    for rebalance_date in dates:
        if date_slices is None:
            slice_ = prepared.loc[prepared[rebalance_col].eq(rebalance_date), [*zscore_columns, forward_return_col]].dropna()
        else:
            slice_ = date_slices.get(pd.Timestamp(rebalance_date), pd.DataFrame())
            if slice_.empty:
                continue
        slice_ = slice_.loc[np.isfinite(slice_[forward_return_col])]
        if len(slice_) < min_assets:
            continue
        score = slice_[zscore_columns].to_numpy(dtype=float) @ weights
        returns = slice_[forward_return_col]
        if objective == "rank_ic":
            value = _rank_ic(pd.Series(score, index=returns.index), returns)
        elif objective == "ic":
            value = _pearson_ic(pd.Series(score, index=returns.index), returns)
        else:
            groups = _quantile_groups(pd.Series(score, index=returns.index), quantiles)
            value = float(returns.loc[groups.eq(quantiles)].mean() - returns.loc[groups.eq(1)].mean())
        if pd.notna(value):
            values.append(float(value))
    return (float(np.mean(values)), len(values)) if values else (float("nan"), 0)


def optimize_weights_walk_forward(
    panel: pd.DataFrame,
    factor_columns: Sequence[str],
    *,
    rebalance_col: str = "rebalance_date",
    security_col: str = "source_code",
    available_date_col: str = "available_date",
    forward_return_col: str = "forward_return",
    train_periods: int = 12,
    objective: str = "rank_ic",
    candidate_count: int = 256,
    max_weight: float = 1.0,
    random_state: int = 7,
    quantiles: int = 5,
    transaction_cost_bps: float = 0.0,
    long_short: bool = True,
    winsor_limits: tuple[float, float] = (0.01, 0.99),
    min_assets: int = 10,
) -> WeightOptimizationResult:
    """Find simplex weights using trailing data, then test only the next period.

    Every row in ``weight_history`` for date *t* was fitted with dates strictly
    before *t*.  Candidate weights are non-negative and sum to one; a fixed
    random seed makes the candidate search reproducible.
    """
    _validate_options(quantiles, winsor_limits, min_assets)
    if train_periods < 1:
        raise ValueError("train_periods 必须至少为 1")
    if objective not in {"rank_ic", "ic", "long_short_return"}:
        raise ValueError("objective 仅支持 rank_ic、ic 或 long_short_return")
    prepared = prepare_factor_panel(
        panel, factor_columns, rebalance_col=rebalance_col, security_col=security_col,
        available_date_col=available_date_col, forward_return_col=forward_return_col,
        winsor_limits=winsor_limits,
    )
    dates = list(pd.Series(prepared[rebalance_col].dropna().unique()).sort_values())
    candidates = _weight_candidates(len(factor_columns), candidate_count, random_state, max_weight)
    zscore_columns = [f"{factor}_zscore" for factor in factor_columns]
    score_columns = [*zscore_columns, forward_return_col]
    date_slices = {
        pd.Timestamp(rebalance_date): slice_[score_columns].dropna()
        for rebalance_date, slice_ in prepared.groupby(rebalance_col, sort=False)
    }
    history_rows: list[dict[str, Any]] = []
    oos_rows: list[pd.DataFrame] = []
    for position in range(train_periods, len(dates)):
        rebalance_date = dates[position]
        train_dates = dates[position - train_periods:position]
        best_weight: np.ndarray | None = None
        best_score = -np.inf
        best_periods = 0
        for weights in candidates:
            score, valid_periods = _score_weights(
                prepared, factor_columns, weights, train_dates, rebalance_col=rebalance_col,
                forward_return_col=forward_return_col, min_assets=min_assets,
                quantiles=quantiles, objective=objective, date_slices=date_slices,
            )
            if valid_periods < train_periods or not np.isfinite(score):
                continue
            if score > best_score:
                best_weight, best_score, best_periods = weights, score, valid_periods
        if best_weight is None:
            continue
        oos = prepared.loc[prepared[rebalance_col].eq(rebalance_date)].copy()
        oos["composite_factor"] = oos[zscore_columns].to_numpy(dtype=float) @ best_weight
        oos_rows.append(oos)
        record: dict[str, Any] = {
            "rebalance_date": rebalance_date,
            "training_start": train_dates[0],
            "training_end": train_dates[-1],
            "training_periods": best_periods,
            "objective": objective,
            "training_objective": best_score,
        }
        record.update({factor: float(weight) for factor, weight in zip(factor_columns, best_weight)})
        history_rows.append(record)
    history_columns = [
        "rebalance_date", "training_start", "training_end", "training_periods", "objective",
        "training_objective", *factor_columns,
    ]
    history = pd.DataFrame(history_rows, columns=history_columns)
    if not oos_rows:
        empty_panel = prepared.iloc[0:0].copy()
        empty_panel["composite_factor"] = pd.Series(dtype=float)
        return WeightOptimizationResult(
            weights=pd.Series(dtype=float, name="weight"),
            weight_history=history,
            oos_panel=empty_panel,
            oos_backtest=_empty_backtest(empty_panel),
        )
    oos_panel = pd.concat(oos_rows, ignore_index=True)
    backtest = run_factor_backtest(
        oos_panel, ["composite_factor"], rebalance_col=rebalance_col, security_col=security_col,
        available_date_col=available_date_col, forward_return_col=forward_return_col,
        quantiles=quantiles, transaction_cost_bps=transaction_cost_bps, long_short=long_short,
        winsor_limits=winsor_limits, min_assets=min_assets,
    )
    latest = history.iloc[-1][list(factor_columns)].astype(float).rename("weight")
    return WeightOptimizationResult(latest, history, oos_panel, backtest)


# Readable aliases for callers that use verb-first naming.
backtest_factors = run_factor_backtest
walk_forward_optimize_weights = optimize_weights_walk_forward