import sys
import time
import traceback

import pandas as pd
import numpy as np

from config import SCREENER_TOP_N
from data.sina_fetcher import sina_get_market_snapshot, sina_get_stock_hist
from data.preprocessor import clean_snapshot, prepare_hist
from features.technical import compute_technical_features, get_latest_features
from features.money_flow import compute_flow_features
from features.market_context import compute_market_context, compute_sector_features
from models.ensemble import EnsembleModel
from llm.analyzer import StockAnalyzer
from llm.client import LLMClient
from strategy.screener import apply_risk_filters, apply_model_risk_gates, rank_stocks
from strategy.backtest import generate_labels


def run_screener(top_n: int = SCREENER_TOP_N, use_llm: bool = True,
                 max_stocks: int = 30, mode: str = "auto"):
    """运行选股流程。
    mode: auto(自动探测) / pre(盘前规划) / confirm(盘中确认)
    """
    print("[1/6] 获取市场数据...")
    snapshot = sina_get_market_snapshot(max_pages=30)
    if snapshot.empty:
        print("  ERROR: 无法获取市场快照")
        return None, {}, ""

    snapshot = clean_snapshot(snapshot)
    print(f"  过滤后剩余 {len(snapshot)} 只股票")

    live_amount = snapshot["成交额(元)"].sum()
    is_premarket = live_amount <= 0
    if mode == "auto":
        mode = "pre" if is_premarket else "confirm"
    elif mode == "pre":
        is_premarket = True
    elif mode == "confirm":
        is_premarket = False

    from data.fetcher import save_snapshot, load_last_snapshot
    pool_source = None
    if is_premarket:
        prev = load_last_snapshot()
        if not prev.empty:
            prev = clean_snapshot(prev)
            pool_source = prev.sort_values("成交额(元)", ascending=False)
            market_ctx = compute_market_context(prev)
            print(f"  盘前规划版: 候选池基于昨日快照(成交额) {len(pool_source)} 只")
        else:
            pool_source = snapshot.sort_values("总市值(元)", ascending=False)
            market_ctx = compute_market_context(snapshot)
            print("  盘前规划版: 无历史快照，回退总市值排序候选池")
    else:
        save_snapshot(snapshot)
        pool_source = snapshot.sort_values("成交额(元)", ascending=False)
        market_ctx = compute_market_context(snapshot)
        print(f"  盘中确认版: 候选池基于实时成交额 {len(pool_source)} 只")

    print("[2/6] 计算市场环境...")
    print(f"  上涨{market_ctx.get('market_advance_count', 0)} 下跌{market_ctx.get('market_decline_count', 0)} "
          f"均涨幅{market_ctx.get('market_mean_return', 0):.2f}%")

    print(f"[3/6] 获取Top{max_stocks}股票K线...")
    top_codes = pool_source.head(max_stocks)["代码"].tolist()

    all_features = []
    for i, code in enumerate(top_codes):
        row = pool_source[pool_source["代码"] == code]
        if row.empty:
            continue
        r = row.iloc[0]
        feat = {
            "code": code,
            "name": r.get("名称", ""),
            "price": r.get("最新价", 0),
            "change_pct": r.get("涨跌幅(%)", 0),
            "amount": r.get("成交额(元)", 0),
            "turnover": r.get("换手率(%)", 0),
        }

        hist = sina_get_stock_hist(code, days=60)
        if not hist.empty:
            hist = prepare_hist(hist)
            tech = get_latest_features(hist)
            feat.update(tech)
            if is_premarket and not hist.empty:
                last = hist.iloc[-1]
                feat["price"] = last.get("close", feat["price"])
                feat["change_pct"] = last.get("change_pct", feat["change_pct"])

        for k, v in market_ctx.items():
            feat[k] = v

        all_features.append(feat)
        if (i + 1) % 10 == 0:
            print(f"  已处理 {i + 1}/{len(top_codes)}")
        time.sleep(0.3)

    features_df = pd.DataFrame(all_features)
    if features_df.empty:
        print("  ERROR: 特征计算失败")
        return None, {}, mode
    print(f"  特征矩阵: {features_df.shape}")

    print("[4/6] ML模型打分...")
    numeric_cols = features_df.select_dtypes(include=[np.number]).columns.tolist()
    exclude_cols = ["price", "change_pct", "amount"]
    all_feature_cols = [c for c in numeric_cols if c not in exclude_cols]

    model = EnsembleModel()
    if model.lgbm.load("ensemble_lgbm") and model.xgb.load("ensemble_xgb"):
        model_features = model.lgbm.feature_names
        aligned_cols = [c for c in model_features if c in all_feature_cols]
        missing = [c for c in model_features if c not in features_df.columns]
        for c in missing:
            features_df[c] = 0
            aligned_cols.append(c) if c in model_features else None
        X = features_df[model_features].fillna(0)
        scores = model.score_stocks(X)
        features_df = pd.concat([features_df, scores], axis=1)
        print(f"  使用训练好的模型打分 (特征对齐: {len(model_features)}个)")
    else:
        X = features_df[all_feature_cols].fillna(0)
        raw_score = X.mean(axis=1)
        features_df["ml_score"] = (raw_score - raw_score.min()) / (raw_score.max() - raw_score.min() + 1e-10)
        features_df["ml_score_norm"] = features_df["ml_score"]
        print("  未找到训练模型，使用特征排名替代")

    print("[5/6] 综合排名...")
    ranked = rank_stocks(features_df)
    ranked = apply_model_risk_gates(ranked, {"min_score": 0.0, "max_rank": top_n})
    ranked = ranked.copy()
    ranked["mode"] = mode
    result = ranked.head(top_n)

    # 添加买入概率和建议
    # 方案B：排名相对强度(p^4幂曲线拉开头部) + 绝对护栏(ml_score过低时封顶)
    pool_size = max(len(features_df), 1)

    def _calibrate_buy_prob(row):
        rank = int(row.get("rank", pool_size))
        ml = row.get("ml_score", 0.0)
        p = max(0.0, 1.0 - (rank - 0.5) / pool_size)
        rel = p ** 4
        base = 0.12 + 0.82 * rel
        if ml < 0.22:
            base = min(base, 0.40)
        elif ml < 0.26:
            base = min(base, 0.50)
        elif ml < 0.32:
            base = min(base, 0.60)
        return max(0.10, min(0.95, base))

    result["buy_prob"] = result.apply(_calibrate_buy_prob, axis=1)
    result["advice"] = result["buy_prob"].apply(_get_advice)

    mode_label = "盘前规划版" if mode == "pre" else "盘中确认版"
    hint = "预测今日走向，结合9:25竞价观察，10:00确认买入" if mode == "pre" else "实时数据，决策买入，次日卖出"

    print("[6/6] 输出结果...")
    print()
    print("=" * 72)
    print(f"  A股超短期选股 {mode_label} Top {top_n}  {hint}")
    print("=" * 72)
    print(f"  {'排名':<4} {'代码':<8} {'名称':<10} {'买入概率':<10} {'涨跌幅':<10} {'建议'}")
    print("-" * 72)
    for _, row in result.iterrows():
        prob = row["buy_prob"]
        ml_raw = row.get("ml_score", 0)
        print(
            f"  {int(row['rank']):<4} {row['code']:<8} {row['name']:<10} "
            f"{prob*100:>6.1f}%    {row.get('change_pct', 0):>+6.2f}%    {row['advice']}"
            f"    (ML原始分: {ml_raw:.3f})"
        )
    print("=" * 72)
    if mode == "pre":
        print("  操作：盘前观察名单，9:25看竞价强弱，10:00用确认版决定买入")
    else:
        print("  操作：从前5名选2-3只买入 | 止损：亏5%卖 | 止盈：赚5%卖")
    print("=" * 72)

    # LLM分析（可选）
    if use_llm:
        client = LLMClient()
        if client.is_available():
            print("\n  LLM分析中...")
            analyzer = StockAnalyzer(client)
            llm_report = analyzer.generate_report(
                result.to_dict("records"), market_ctx
            )
            print(llm_report)
        else:
            print("\n  LLM未配置或不可用，跳过LLM分析")

    return result, market_ctx, mode


def _find_col(df, candidates):
    for c in candidates:
        if c in df.columns:
            return c
    return ""


def _get_advice(prob):
    if prob >= 0.65:
        return "强烈推荐"
    elif prob >= 0.55:
        return "可以买入"
    elif prob >= 0.45:
        return "观望"
    else:
        return "不建议"


if __name__ == "__main__":
    result, _, _ = run_screener()
    if result is not None:
        print("\n选股完成!")