"""智能选股引擎 — 可视化选股器、策略保存、回测验证。

功能：
1. 多条件组合选股
2. 选股策略保存/加载
3. 选股结果历史回测
4. 选股结果对比（新入选/退出）
5. 条件预警
"""
import datetime as dt
import json
from typing import Dict, List, Any, Optional
from sqlalchemy import select, and_, or_, func
import numpy as np

from db import get_session
from models import StockMetric, DailyQuote, Security


class Condition:
    """选股条件"""
    
    def __init__(self, field: str, operator: str, value: Any):
        self.field = field
        self.operator = operator
        self.value = value
    
    def to_dict(self):
        return {
            "field": self.field,
            "operator": self.operator,
            "value": self.value
        }
    
    @classmethod
    def from_dict(cls, data: Dict):
        return cls(data["field"], data["operator"], data["value"])
    
    def to_sql(self, model):
        """转换为 SQLAlchemy 查询条件"""
        field = getattr(model, self.field, None)
        if field is None:
            return None
        
        op = self.operator
        val = self.value
        
        if op == "==":
            return field == val
        elif op == ">":
            return field > val
        elif op == ">=":
            return field >= val
        elif op == "<":
            return field < val
        elif op == "<=":
            return field <= val
        elif op == "between":
            return and_(field >= val[0], field <= val[1])
        elif op == "in":
            return field.in_(val)
        else:
            return None


class Strategy:
    """选股策略"""
    
    def __init__(self, name: str, description: str = ""):
        self.name = name
        self.description = description
        self.conditions: List[Condition] = []
        self.logic = "and"  # and / or
    
    def add_condition(self, field: str, operator: str, value: Any):
        """添加条件"""
        self.conditions.append(Condition(field, operator, value))
    
    def to_dict(self):
        return {
            "name": self.name,
            "description": self.description,
            "logic": self.logic,
            "conditions": [c.to_dict() for c in self.conditions]
        }
    
    @classmethod
    def from_dict(cls, data: Dict):
        strategy = cls(data["name"], data.get("description", ""))
        strategy.logic = data.get("logic", "and")
        for cond in data.get("conditions", []):
            strategy.conditions.append(Condition.from_dict(cond))
        return strategy
    
    def to_json(self) -> str:
        """序列化为JSON"""
        return json.dumps(self.to_dict(), ensure_ascii=False)
    
    @classmethod
    def from_json(cls, json_str: str):
        """从JSON反序列化"""
        return cls.from_dict(json.loads(json_str))


def run_selector(strategy: Strategy, date: Optional[dt.date] = None) -> Dict[str, Any]:
    """执行选股
    
    Args:
        strategy: 选股策略
        date: 选股日期，None表示最新
    
    Returns:
        选股结果
    """
    with get_session() as s:
        # 确定日期
        if date is None:
            date = s.execute(select(func.max(StockMetric.date))).scalar()
        
        if not date:
            return {"ok": False, "msg": "暂无数据"}
        
        # 构建查询
        query = select(StockMetric).where(StockMetric.date == date)
        
        # 应用条件
        sql_conditions = []
        for cond in strategy.conditions:
            sql_cond = cond.to_sql(StockMetric)
            if sql_cond is not None:
                sql_conditions.append(sql_cond)
        
        if sql_conditions:
            if strategy.logic == "and":
                query = query.where(and_(*sql_conditions))
            else:  # or
                query = query.where(or_(*sql_conditions))
        
        # 执行查询
        rows = s.execute(query).scalars().all()
        
        # 格式化结果
        results = [{
            "code": r.code,
            "name": r.name,
            "close": round(r.close, 2),
            "pct": round(r.pct, 2),
            "ret5": round(r.ret5, 2),
            "ret20": round(r.ret20, 2),
            "vol_ratio": round(r.vol_ratio, 2),
            "rsi14": round(r.rsi14, 2),
            "pos60": round(r.pos60 * 100, 1),
            "amount": round(r.amount, 2),
            "ma_bull": r.ma_bull,
            "macd_gold": r.macd_gold
        } for r in rows]
    
    return {
        "ok": True,
        "date": date.isoformat(),
        "strategy": strategy.name,
        "count": len(results),
        "results": results
    }


def backtest_selector(strategy: Strategy, days: int = 60) -> Dict[str, Any]:
    """选股策略回测
    
    Args:
        strategy: 选股策略
        days: 回测天数
    
    Returns:
        回测结果
    """
    with get_session() as s:
        # 获取最近N个交易日
        latest_date = s.execute(select(func.max(StockMetric.date))).scalar()
        if not latest_date:
            return {"ok": False, "msg": "暂无数据"}
        
        start_date = latest_date - dt.timedelta(days=days)
        
        # 获取这段时间内的所有交易日
        dates = s.execute(
            select(StockMetric.date)
            .where(StockMetric.date >= start_date)
            .group_by(StockMetric.date)
            .order_by(StockMetric.date)
        ).scalars().all()
    
    if len(dates) < 5:
        return {"ok": False, "msg": "数据不足"}
    
    # 逐日选股并统计后续N日收益
    daily_results = []
    for i, date in enumerate(dates[:-5]):  # 至少保留5日用于计算收益
        # 执行选股
        result = run_selector(strategy, date)
        if not result["ok"] or not result["results"]:
            continue
        
        selected_codes = [r["code"] for r in result["results"]]
        
        # 查询5日后的收益
        future_date = dates[min(i + 5, len(dates) - 1)]
        
        with get_session() as s:
            # 获取选中股票5日后的表现
            future_rows = s.execute(
                select(DailyQuote.code, DailyQuote.close)
                .where(
                    and_(
                        DailyQuote.code.in_(selected_codes),
                        DailyQuote.date == future_date
                    )
                )
            ).all()
            
            current_rows = s.execute(
                select(DailyQuote.code, DailyQuote.close)
                .where(
                    and_(
                        DailyQuote.code.in_(selected_codes),
                        DailyQuote.date == date
                    )
                )
            ).all()
        
        # 计算收益
        current_prices = {r.code: float(r.close) for r in current_rows}
        future_prices = {r.code: float(r.close) for r in future_rows}
        
        returns = []
        for code in selected_codes:
            if code in current_prices and code in future_prices:
                ret = (future_prices[code] / current_prices[code] - 1) * 100
                returns.append(ret)
        
        if returns:
            daily_results.append({
                "date": date.isoformat(),
                "count": len(selected_codes),
                "avg_return": round(np.mean(returns), 2),
                "median_return": round(np.median(returns), 2),
                "win_rate": round(sum(1 for r in returns if r > 0) / len(returns) * 100, 1),
                "max_return": round(max(returns), 2),
                "min_return": round(min(returns), 2)
            })
    
    if not daily_results:
        return {"ok": False, "msg": "回测数据不足"}
    
    # 汇总统计
    avg_returns = [r["avg_return"] for r in daily_results]
    win_rates = [r["win_rate"] for r in daily_results]
    
    summary = {
        "total_days": len(daily_results),
        "avg_return": round(np.mean(avg_returns), 2),
        "avg_win_rate": round(np.mean(win_rates), 1),
        "best_day": max(daily_results, key=lambda x: x["avg_return"]),
        "worst_day": min(daily_results, key=lambda x: x["avg_return"])
    }
    
    return {
        "ok": True,
        "strategy": strategy.name,
        "days": days,
        "summary": summary,
        "daily": daily_results
    }


def compare_results(date1: dt.date, date2: dt.date, strategy: Strategy) -> Dict[str, Any]:
    """对比两个日期的选股结果
    
    Args:
        date1: 日期1（通常是昨日）
        date2: 日期2（通常是今日）
        strategy: 选股策略
    
    Returns:
        对比结果
    """
    result1 = run_selector(strategy, date1)
    result2 = run_selector(strategy, date2)
    
    if not result1["ok"] or not result2["ok"]:
        return {"ok": False, "msg": "数据不足"}
    
    codes1 = set(r["code"] for r in result1["results"])
    codes2 = set(r["code"] for r in result2["results"])
    
    # 新入选
    new_in = codes2 - codes1
    # 退出
    dropped = codes1 - codes2
    # 持续入选
    continued = codes1 & codes2
    
    # 获取详细信息
    new_in_stocks = [r for r in result2["results"] if r["code"] in new_in]
    dropped_stocks = [r for r in result1["results"] if r["code"] in dropped]
    continued_stocks = [r for r in result2["results"] if r["code"] in continued]
    
    return {
        "ok": True,
        "date1": date1.isoformat(),
        "date2": date2.isoformat(),
        "count1": len(codes1),
        "count2": len(codes2),
        "new_in": {
            "count": len(new_in),
            "stocks": new_in_stocks
        },
        "dropped": {
            "count": len(dropped),
            "stocks": dropped_stocks
        },
        "continued": {
            "count": len(continued),
            "stocks": continued_stocks
        }
    }


# 预设策略
PRESET_STRATEGIES = {
    "momentum": Strategy("动量突破", "短期强势+放量"),
    "value": Strategy("价值洼地", "超跌低位+基本面支撑"),
    "growth": Strategy("成长加速", "持续上涨+量价齐升"),
    "reversal": Strategy("反转抄底", "超跌企稳+技术反转")
}

# 动量突破
PRESET_STRATEGIES["momentum"].add_condition("ret5", ">", 10)
PRESET_STRATEGIES["momentum"].add_condition("vol_ratio", ">", 2)
PRESET_STRATEGIES["momentum"].add_condition("rsi14", "<", 80)

# 价值洼地
PRESET_STRATEGIES["value"].add_condition("pos60", "<", 0.3)
PRESET_STRATEGIES["value"].add_condition("pct", ">", 0)
PRESET_STRATEGIES["value"].add_condition("amount", ">", 5)

# 成长加速
PRESET_STRATEGIES["growth"].add_condition("ret20", ">", 15)
PRESET_STRATEGIES["growth"].add_condition("ma_bull", "==", True)
PRESET_STRATEGIES["growth"].add_condition("up_streak", ">=", 2)

# 反转抄底
PRESET_STRATEGIES["reversal"].add_condition("ret20", "<", -15)
PRESET_STRATEGIES["reversal"].add_condition("rsi14", "<", 30)
PRESET_STRATEGIES["reversal"].add_condition("pct", ">", 2)


def get_preset_strategies() -> List[Dict[str, Any]]:
    """获取预设策略列表"""
    return [
        {
            "id": key,
            "name": strategy.name,
            "description": strategy.description,
            "conditions_count": len(strategy.conditions)
        }
        for key, strategy in PRESET_STRATEGIES.items()
    ]


def get_available_fields() -> List[Dict[str, Any]]:
    """获取可用的选股字段"""
    return [
        {"field": "close", "name": "现价", "type": "number", "unit": "元"},
        {"field": "pct", "name": "涨跌幅", "type": "number", "unit": "%"},
        {"field": "ret5", "name": "5日涨幅", "type": "number", "unit": "%"},
        {"field": "ret20", "name": "20日涨幅", "type": "number", "unit": "%"},
        {"field": "ret60", "name": "60日涨幅", "type": "number", "unit": "%"},
        {"field": "ma5", "name": "MA5", "type": "number", "unit": "元"},
        {"field": "ma10", "name": "MA10", "type": "number", "unit": "元"},
        {"field": "ma20", "name": "MA20", "type": "number", "unit": "元"},
        {"field": "ma60", "name": "MA60", "type": "number", "unit": "元"},
        {"field": "vol_ratio", "name": "量比", "type": "number", "unit": ""},
        {"field": "rsi14", "name": "RSI", "type": "number", "unit": ""},
        {"field": "pos60", "name": "60日分位", "type": "number", "unit": ""},
        {"field": "amount", "name": "成交额", "type": "number", "unit": "亿"},
        {"field": "up_streak", "name": "连涨天数", "type": "number", "unit": "天"},
        {"field": "ma_bull", "name": "均线多头", "type": "boolean", "unit": ""},
        {"field": "macd_gold", "name": "MACD金叉", "type": "boolean", "unit": ""}
    ]