# 克隆自聚宽文章：https://www.joinquant.com/post/67602
# 标题：ETF动量策略优化
# 作者：金大炮

# ==================== ETF轮动策略 · 原版加温和止损 + 动态核心池优化 + 资金热度因子 ====================
# 优化功能：
#   1. 动态核心池：每季度从备选池中根据长期动量筛选前20只
#   2. 保留动态热点池：每周更新成交额最高的行业ETF
#   3. 最终候选池 = 动态核心池 + 动态热点池
#   4. 保留原版12%硬止损
#   5. 新增资金热度因子（近5日/近20日成交额变化），与动量得分加权
# ================================================================================================

import math
import numpy as np
import pandas as pd
from jqdata import *

def initialize(context):
    set_option("avoid_future_data", True)
    set_option("use_real_price", True)
    
    print("【初始化】策略开始运行，当前日期：", context.current_dt)
    
    log.set_level('order', 'error')
    log.set_level('system', 'error')
    log.set_level('strategy', 'debug')
    
    # 滑点与佣金（实盘级）
    set_slippage(FixedSlippage(0.0005), type="fund")
    set_slippage(FixedSlippage(0.001), type="stock")
    
    set_order_cost(
        OrderCost(
            open_tax=0, close_tax=0.001,
            open_commission=0.0003, close_commission=0.0003,
            close_today_commission=0, min_commission=5
        ),
        type="stock"
    )
    
    set_order_cost(
        OrderCost(
            open_tax=0, close_tax=0,
            open_commission=0.0003, close_commission=0.0003,
            close_today_commission=0, min_commission=0
        ),
        type="mmf"
    )

    # 全局变量
    g.strategys = {}
    g.portfolio_value_proportion = [1]
    g.positions = {i: {} for i in range(len(g.portfolio_value_proportion))}
    g.dynamic_etf_pool = []          # 动态热点行业池
    g.weights = {}
    g.position_ratio = 1.0            # 始终满仓

    # 策略参数
    g.m_days = 45                     # 动量周期
    g.stock_sum =5                   # 持仓ETF数量
    g.hot_money_days = 20              # 行业池成交额天数
    g.stop_loss_pct = 0.12             # 硬止损幅度（12%）

    # ========== 新增：动态核心池相关 ==========
    g.long_term_momentum_days = 250     # 长期动量周期（用于筛选核心池）
    g.core_pool_update_freq = 90         # 每90天更新一次核心池
    g.last_core_pool_update = None       # 上次更新时间
    g.dynamic_core_pool = []             # 动态核心池
    
    # ========== 新增：资金热度因子相关 ==========
    g.money_heat_short_days = 5          # 短期成交额平均天数
    g.money_heat_long_days = 30           # 长期成交额平均天数
    g.momentum_weight = 0.7               # 动量权重
    g.money_heat_weight = 0.3             # 资金热度权重
    
    # 建立备选核心池（约50只，覆盖各类优质ETF）
    g.candidate_core_pool = [
        # 宽基指数
        "510050.XSHG",  # 上证50
        "510300.XSHG",  # 沪深300
        "510500.XSHG",  # 中证500
        "159915.XSHE",  # 创业板
        "588000.XSHG",  # 科创50
        "159949.XSHE",  # 创业板50
        "512100.XSHG",  # 中证1000
        "588080.XSHG",  # 科创50（易方达）
        "159592.XSHE",  # A50
        "512160.XSHG",  # MSCI中国A股
        # 行业/主题
        "512880.XSHG",  # 证券
        "512660.XSHG",  # 军工
        "515790.XSHG",  # 光伏
        "512400.XSHG",  # 有色金属
        "159825.XSHE",  # 农业
        "515030.XSHG",  # 新能源车
        "512170.XSHG",  # 医疗
        "159928.XSHE",  # 消费
        "515050.XSHG",  # 5G
        "512480.XSHG",  # 半导体
        "512760.XSHG",  # 芯片
        "515700.XSHG",  # 新能源
        "159869.XSHE",  # 传媒
        "512690.XSHG",  # 酒
        "159766.XSHE",  # 旅游
        "516160.XSHG",  # 新能源
        "516550.XSHG",  # 大数据
        "517050.XSHG",  # 互联网
        # 商品/跨境
        "518880.XSHG",  # 黄金
        "159985.XSHE",  # 豆粕
        "513100.XSHG",  # 纳指
        "513520.XSHG",  # 日经
        "513030.XSHG",  # 德国
        "159980.XSHE",  # 有色
        "501018.XSHG",  # 南方原油
        "513130.XSHG",  # 恒生科技
        "164906.XSHE",  # 中国互联
        "161125.XSHE",  # 标普500
        "513080.XSHG",  # 法国CAC40
        "513090.XSHG",  # 香港证券
        "513550.XSHG",  # 港股通50
        "513660.XSHG",  # 恒生国企
        "159937.XSHE",  # 博时黄金
        "159934.XSHE",  # 黄金ETF
        # 赛道类（原有的一些）
        "512290.XSHG",  # 生物医药
        "515070.XSHG",  # AI智能
        "159851.XSHE",  # 金融科技
        "159637.XSHE",  # 新能源车
        "159550.XSHE",  # 互联网
        "512710.XSHG",  # 军工龙头
        "159692.XSHE",  # 证券ETF东财
    ]

    if g.portfolio_value_proportion[0] > 0:
        run_weekly(etf_rotation_adjust, 4, "14:50")   # 周五调仓
    
    run_daily(end_trade, "14:59")
    run_weekly(update_sector_pool, 0, "09:00")       # 周一更新行业池
    run_monthly(update_core_pool, 1, "09:30")        # 每月第一个交易日更新核心池

    process_initialize(context)


def update_core_pool(context):
    """每季度更新动态核心池（基于长期动量）"""
    print("【核心池更新】开始执行...")
    
    # 检查是否需要更新
    if g.last_core_pool_update is not None:
        days_since_update = (context.current_dt.date() - g.last_core_pool_update).days
        if days_since_update < g.core_pool_update_freq:
            print(f"【核心池】距离上次更新仅{days_since_update}天，暂不更新")
            return
    
    print("【核心池】开始计算长期动量，筛选核心ETF...")
    end_date = context.previous_date
    
    # 计算备选池中每个ETF的长期动量
    momentum_scores = []
    for code in g.candidate_core_pool:
        try:
            # 获取足够的历史数据
            df = attribute_history(code, g.long_term_momentum_days, '1d', ['close'])
            if df is None or len(df) < g.long_term_momentum_days:
                continue
            
            # 计算简单长期动量（百分比）
            start_price = df['close'].iloc[0]
            end_price = df['close'].iloc[-1]
            momentum = (end_price / start_price - 1) * 100
            
            # 可选：加入波动率过滤（波动太大的剔除）
            returns = df['close'].pct_change().dropna()
            volatility = returns.std() * np.sqrt(252) * 100  # 年化波动率
            
            # 综合得分：动量 - 波动率惩罚（可选）
            # score = momentum - volatility * 0.5
            score = momentum  # 简单版本只用动量
            
            momentum_scores.append((code, score, momentum, volatility))
            
        except Exception as e:
            print(f"【核心池】计算{code}动量出错: {e}")
            continue
    
    if not momentum_scores:
        print("【核心池】无有效数据，保持原池")
        return
    
    # 按得分排序
    momentum_scores.sort(key=lambda x: x[1], reverse=True)
    
    # 选出前20名作为新的核心池
    top_count = min(20, len(momentum_scores))
    new_core_pool = [code for code, score, mom, vol in momentum_scores[:top_count]]
    
    # 更新全局变量
    g.dynamic_core_pool = new_core_pool
    g.last_core_pool_update = context.current_dt.date()
    
    print(f"【核心池更新完成】新核心池 ({len(new_core_pool)}只):")
    for i, code in enumerate(new_core_pool[:10]):  # 只打印前10只
        name = get_security_info(code).display_name
        print(f"  {i+1}. {name} ({code})")
    if len(new_core_pool) > 10:
        print(f"  ... 等共{len(new_core_pool)}只")


def update_sector_pool(context):
    """【增强版】动态行业池 - 增加质量过滤和数量"""
    print("【行业池更新·增强版】正在执行...")
    
    # 获取所有ETF
    all_etfs = get_all_securities(['etf']).index.tolist()
    
    # 排除词
    exclude_keywords = ['300', '500', '1000', '50', '货币', '债', '国债', '地方债']
    
    # 第一步：初筛（排除明显不想投的）
    candidate_etfs = []
    current_data = get_current_data()
    
    for code in all_etfs:
        try:
            name = get_security_info(code).display_name
            
            # 排除包含关键词的
            should_exclude = False
            for k in exclude_keywords:
                if k in name:
                    should_exclude = True
                    break
            if should_exclude:
                continue
            
            # 质量过滤：日均成交额 > 500万，价格 > 0.5，非停牌
            if code in current_data:
                if current_data[code].paused:
                    continue
                if current_data[code].last_price < 0.5:
                    continue
            
            candidate_etfs.append(code)
            
        except Exception as e:
            continue
    
    if not candidate_etfs:
        print("未找到符合条件的行业ETF")
        return
    
    # 第二步：计算过去20日平均成交额
    end_date = context.previous_date
    try:
        h = get_price(candidate_etfs, count=g.hot_money_days, end_date=end_date, 
                      frequency='daily', fields=['money'], skip_paused=True)
        avg_money = h['money'].mean().sort_values(ascending=False)
    except Exception as e:
        print(f"获取成交额数据出错: {e}")
        return
    
    # 第三步：选出成交额前20名
    top_codes = avg_money.head(20).index.tolist()
    
    # 第四步：行业去重（取每个行业成交额最高的）
    final_dynamic_pool = []
    seen_industries = set()
    
    for code in top_codes:
        try:
            name = get_security_info(code).display_name
            # 简单的行业判断（取前2个汉字）
            industry_key = name[:2]
            
            if industry_key not in seen_industries:
                final_dynamic_pool.append(code)
                seen_industries.add(industry_key)
                
            if len(final_dynamic_pool) >= 8:  # 增加到8只
                break
        except:
            continue
    
    # 如果不足8只，补充一些未去重的（但保证成交额高）
    if len(final_dynamic_pool) < 8:
        for code in top_codes:
            if code not in final_dynamic_pool and len(final_dynamic_pool) < 8:
                final_dynamic_pool.append(code)
    
    g.dynamic_etf_pool = final_dynamic_pool
    print(f"【动态更新·增强版】本周热点行业池 ({len(final_dynamic_pool)}只):")
    for code in final_dynamic_pool:
        print(f"  {get_security_info(code).display_name} ({code})")


def end_trade(context):
    marked = {s for d in g.positions.values() for s in d}
    current_data = get_current_data()
    for stock in context.portfolio.positions:
        if stock not in marked:
            price = current_data[stock].last_price
            pos = context.portfolio.positions[stock].total_amount
            if my_order(stock, -pos, price, 0):
                print(f"卖出{stock}因送股未记录在持仓中", price, pos)


def my_order(security, vol, price, target_position):
    o = order(security, vol)
    return o


def etf_rotation_adjust(context):
    print("【调仓】etf_rotation_adjust 被触发，时间：", context.current_dt)
    g.strategys["核心资产轮动策略"].adjust()


def process_initialize(context):
    print("重启程序")
    g.strategys = {
        name: cls(context, index=idx, name=name)
        for name, cls, idx in [
            ("核心资产轮动策略", Etf_Rotation_Strategy, 0),
        ]
    }


# -------------------- 策略基类 --------------------
class Strategy:
    def __init__(self, context, index, name):
        self.context = context
        self.index = index
        self.name = name
        self.stock_sum = g.stock_sum
        self.hold_list = []
        self.min_money = 500

    def get_total_value(self):
        if not g.positions[self.index]:
            return 0
        return sum(self.context.portfolio.positions[key].price * value for key, value in g.positions[self.index].items())

    def _adjust(self, targets):
        current_data = get_current_data()
        self.hold_list = list(g.positions[self.index].keys())
        portfolio = self.context.portfolio
        target_value = self.context.portfolio.total_value * g.portfolio_value_proportion[self.index]
        
        # 卖出被调出的
        for stock in self.hold_list:
            if stock not in targets:
                self.order_target_value_(stock, 0)
        # 先卖后买
        for stock, weight in targets.items():
            target = target_value * weight
            price = current_data[stock].last_price
            value = g.positions[self.index].get(stock, 0) * price
            if value - target > max(self.min_money, price * 100):
                self.order_target_value_(stock, target)
        for stock, weight in targets.items():
            target = target_value * weight
            price = current_data[stock].last_price
            value = g.positions[self.index].get(stock, 0) * price
            if min(target - value, portfolio.available_cash) > max(self.min_money, price * 100):
                self.order_target_value_(stock, target)

    def order_target_value_(self, security, value):
        current_data = get_current_data()
        if current_data[security].paused:
            print(f"{security}: 今日停牌")
            return False
        if current_data[security].last_price == current_data[security].high_limit:
            print(f"{security}: 当前涨停")
            return False
        if current_data[security].last_price == current_data[security].low_limit:
            print(f"{security}: 当前跌停")
            return False
        price = current_data[security].last_price
        current_position = g.positions[self.index].get(security, 0)
        current_position_all = self.context.portfolio.positions[security].total_amount if security in self.context.portfolio.positions else 0
        target_position = (int(value / price) // 100) * 100 if price != 0 else 0
        adjustment = target_position - current_position
        target_position_all = current_position_all + adjustment
        closeable_amount = self.context.portfolio.positions[security].closeable_amount if security in self.context.portfolio.positions else 0
        if adjustment < 0 and closeable_amount == 0:
            print(f"{security}: 当天买入不可卖出")
            return False
        if adjustment != 0:
            o = my_order(security, adjustment, price, target_position_all)
            if o:
                filled = o.filled if o.is_buy else -o.filled
                g.positions[self.index][security] = filled + current_position
                if g.positions[self.index][security] == 0:
                    g.positions[self.index].pop(security, None)
                self.hold_list = list(g.positions[self.index].keys())
                return True
        return False

    def filter_untradeable_stock(self, stocks):
        current_data = get_current_data()
        return [
            stock
            for stock in stocks
            if current_data[stock].paused or current_data[stock].last_price in (current_data[stock].high_limit, current_data[stock].low_limit)
        ]


# -------------------- 核心策略类（加止损 + 动态核心池 + 资金热度因子）--------------------
class Etf_Rotation_Strategy(Strategy):
    def __init__(self, context, index, name):
        super().__init__(context, index, name)
        # 保留原来的固定池作为保底（可选）
        self.etf_pool = [
            "513100.XSHG",  # 纳指ETF
            "513520.XSHG",  # 日经ETF
            "513030.XSHG",  # 德国ETF
            "518880.XSHG",  # 黄金ETF
            "159980.XSHE",  # 有色ETF
            "159985.XSHE",  # 豆粕ETF
            "501018.XSHG",  # 南方原油
            "513130.XSHG",  # 恒生科技
            "510180.XSHG",  # 180ETF
            "159915.XSHE",  # 创业板ETF易方达
            "588120.XSHG",  # 科创100
        ]
        self.etf_pool_2 = [
            "512290.XSHG",  # 生物医药
            "515070.XSHG",  # AI智能
            "159851.XSHE",  # 金融科技
            "159637.XSHE",  # 新能源车
            "159550.XSHE",  # 互联网
            "512710.XSHG",  # 军工龙头
            "159692.XSHE"   # 证券ETF东财
        ]
        self.m_days = g.m_days
        self.scores = None
        # 记录每个ETF的近期高点（用于止损）
        self.peak_prices = {}   # {code: peak_price}

    # ---------- 动量计算 + 资金热度因子 ----------
    def filter(self):
        current_date = self.context.current_dt
        valid_etfs = []
        
        # 合并候选池：动态核心池 + 动态热点池
        candidate_pool = []
        if hasattr(g, 'dynamic_core_pool') and g.dynamic_core_pool:
            candidate_pool.extend(g.dynamic_core_pool)
        if g.dynamic_etf_pool:
            candidate_pool.extend(g.dynamic_etf_pool)
        # 如果动态池为空，回退到固定池（保底）
        if not candidate_pool:
            candidate_pool = self.etf_pool + self.etf_pool_2
        # 去重
        candidate_pool = list(set(candidate_pool))
        
        # 筛选出上市时间足够的ETF
        for code in candidate_pool:
            try:
                sec = get_security_info(code)
                if sec is None:
                    continue
                if sec.start_date <= current_date.date():
                    valid_etfs.append(code)
            except Exception as e:
                print(f"【错误】处理 {code} 时发生异常: {e}")
        
        if not valid_etfs:
            print("【filter】当前无有效ETF")
            return []
        
        data = []  # 改用列表存储结果，便于添加新因子
        current_data = get_current_data()
        
        for etf in valid_etfs:
            try:
                # 获取足够的历史数据（至少需要 m_days 天，以及资金热度的长期天数）
                need_days = max(self.m_days, g.money_heat_long_days)
                df = attribute_history(etf, need_days, "1d", ["close", "high", "money"], skip_paused=True)
                if len(df) < need_days:
                    continue
                
                # ---------- 1. 计算动量得分（与原版一致）----------
                prices = np.append(df["close"].values, current_data[etf].last_price)
                y = np.log(prices)
                x = np.arange(len(y))
                weights = np.linspace(1, 2, len(y))
                slope, intercept = np.polyfit(x, y, 1, w=weights)
                
                annual_ret = math.exp(slope * 250) - 1
                ss_res = np.sum(weights * (y - (slope * x + intercept)) ** 2)
                ss_tot = np.sum(weights * (y - np.mean(y)) ** 2)
                r2 = 1 - ss_res / ss_tot if ss_tot else 0
                momentum_score = annual_ret * r2
                
                # 近5日新低惩罚（减半）
                recent_high = df['high'].iloc[-5:].max()
                if current_data[etf].last_price < recent_high * 0.95:
                    momentum_score *= 0.5
                
                # ---------- 2. 计算资金热度因子 ----------
                # 取最近短期和长期的平均成交额（注意使用历史数据，不包含当天可能不完整的成交额）
                # 如果数据不足，跳过
                if len(df) >= g.money_heat_long_days:
                    avg_money_short = df['money'].iloc[-g.money_heat_short_days:].mean()
                    avg_money_long = df['money'].iloc[-g.money_heat_long_days:].mean()
                    if avg_money_long > 0:
                        money_heat = avg_money_short / avg_money_long - 1
                        # 截断到合理范围，避免极端值影响
                        money_heat = max(min(money_heat, 0.5), -0.5)
                    else:
                        money_heat = 0
                else:
                    money_heat = 0
                
                # ---------- 3. 综合得分 ----------
                composite_score = g.momentum_weight * momentum_score + g.money_heat_weight * money_heat
                
                # 只保留综合得分为正的ETF
                if composite_score > 0:
                    data.append((etf, composite_score, momentum_score, money_heat))
                    
            except Exception as e:
                print(f"【错误】计算 {etf} 时出错: {e}")
                continue
        
        # 按综合得分排序
        data.sort(key=lambda x: x[1], reverse=True)
        self.scores = [(code, score) for code, score, _, _ in data]  # 存储综合得分供后续使用
        
        print(f"【filter】有效ETF数：{len(data)}，得分>0：{len(data)}")
        if len(data) > 0:
            print(f"【filter】前3综合得分：{[round(x[1],4) for x in data[:3]]}")
        return [x[0] for x in data]

    # ---------- 止损检查 ----------
    def check_stop_loss(self):
        """检查持仓中是否有触发止损的ETF，返回应卖出的列表"""
        if g.stop_loss_pct <= 0:
            return []
        current_data = get_current_data()
        to_sell = []
        for code in list(g.positions[self.index].keys()):
            current_price = current_data[code].last_price
            # 获取记录的峰值
            peak = self.peak_prices.get(code, current_price)
            if current_price > peak:
                self.peak_prices[code] = current_price
                peak = current_price
            # 计算回撤
            drawdown = (peak - current_price) / peak
            if drawdown >= g.stop_loss_pct:
                print(f"【止损】{code} 回撤 {drawdown:.2%}，触发卖出")
                to_sell.append(code)
                # 移除峰值记录
                self.peak_prices.pop(code, None)
        return to_sell

    def adjust(self):
        print("【adjust】开始调仓")
        
        # 1. 先执行止损
        to_sell = self.check_stop_loss()
        for code in to_sell:
            self.order_target_value_(code, 0)
        
        # 2. 计算动量
        sorted_list = self.filter()
        if len(sorted_list) == 0:
            print("【adjust】无符合条件的ETF，本次不调仓")
            return
        
        # 3. 动量断层保护（1.5倍）
        if len(sorted_list) >= 2:
            score1 = self.scores[0][1]   # 注意self.scores现在是列表
            score2 = self.scores[1][1]
            if score1 < score2 * 1.5:
                targets = sorted_list[:2]
            else:
                targets = sorted_list[:self.stock_sum]
        else:
            targets = sorted_list[:self.stock_sum]
        
        # 4. 等权重分配
        weight = round(1.0 / len(targets), 4)
        target_dict = {etf: weight for etf in targets}
        
        # 5. 更新新买入ETF的峰值
        current_data = get_current_data()
        for etf in targets:
            if etf not in self.peak_prices:
                self.peak_prices[etf] = current_data[etf].last_price
        
        # 6. 执行调仓
        self._adjust(target_dict)
        
        print(f"【调仓完成】买入: {[get_security_info(etf).display_name for etf in targets]}")