"""
WQC Chapter2 速通脚本
生成20条Delay1-USA因子 -> 体检 -> 返回第一条五项全绿因子
"""
import pandas as pd
import numpy as np
from wq_core import alpha, get_univ   # 网页Notebook已内置

# ---------- 参数 ----------
REGION   = 'USA'
DELAY    = 1
UNIV     = get_univ('TOP3000')   # 教学关默认宇宙
N_FACTOR = 20
METRICS  = ['fitness','sharpe','turnover','max_weight','sub_univ_ok']
GREEN    = {'fitness':1.0, 'sharpe':1.25, 'turnover':(1,70), 'max_weight':10, 'sub_univ_ok':True}

# ---------- 1. 因子定义池 ----------
raw_expr = {
    'mom5'   : 'rank(ts_return(close,5))',
    'mom10'  : 'rank(ts_return(close,10))',
    'rev5'   : 'rank(-ts_return(close,5))',
    'rev10'  : 'rank(-ts_return(close,10))',
    'vol10'  : 'rank(-ts_std(return,10))',
    'vol20'  : 'rank(-ts_std(return,20))',
    'amt5'   : 'rank(-ts_mean(amount,5))',      # 短期缩量
    'amt10'  : 'rank(-ts_mean(amount,10))',
    'illiq10': 'rank(ts_mean(abs(return)/amount,10))',  # 非流动
    'illiq20': 'rank(ts_mean(abs(return)/amount,20))',
    # 以下再加半衰变体
    'mom5_d7': 'rank(ts_decay_linear(ts_return(close,5),7))',
    'mom5_d10':'rank(ts_decay_linear(ts_return(close,5),10))',
    'rev5_d7': 'rank(-ts_decay_linear(ts_return(close,5),7))',
    'vol10_d7':'rank(-ts_decay_linear(ts_std(return,10),7))',
    'amt5_d7': 'rank(-ts_decay_linear(ts_mean(amount,5),7))',
    'illiq10_d7':'rank(ts_decay_linear(ts_mean(abs(return)/amount,10),7))',
    # 反转+波动组合
    'rev_vol':'rank(-ts_return(close,5)*ts_std(return,10))',
    'mom_vol':'rank(ts_return(close,5)/ts_std(return,10))',
    'amt_rev':'rank(-ts_mean(amount,5)*ts_return(close,5))',
}

# ---------- 2. 批量生成+体检 ----------
result = []
for name, expr in raw_expr.items():
    fac = alpha(expr, UNIV, delay=DELAY)
    rpt = fac.test(region=REGION, delay=DELAY)   # 返回dict
    rpt['name'] = name
    rpt['expr'] = expr
    result.append(rpt)

df = pd.DataFrame(result)

# ---------- 3. 过滤全绿 ----------
def is_green(row):
    ok = (row.fitness >= GREEN['fitness'] and
          row.sharpe   >= GREEN['sharpe'] and
          GREEN['turnover'][0] <= row.turnover <= GREEN['turnover'][1] and
          row.max_weight <= GREEN['max_weight'] and
          row.sub_univ_ok == GREEN['sub_univ_ok'])
    return ok

greens = df[df.apply(is_green, axis=1)].reset_index(drop=True)

# ---------- 4. 输出 ----------
if greens.empty:
    print('>>> 暂无全绿因子，尝试调半衰期或再中性化 <<<')
else:
    pick = greens.iloc[0]          # 首条即可提交
    print('>>> 发现全绿因子！直接复制下方代码去提交页 <<<')
    print(f'# 因子名: {pick.name}')
    print(f'expression = "{pick.expr}"')
    print(f'# Fitness={pick.fitness:.2f} Sharpe={pick.sharpe:.2f} '
          f'Turnover={pick.turnover:.1f}% MaxWeight={pick.max_weight:.1f}%')