# -*- coding: utf-8 -*-
"""B组挖掘 v4：相对分位果 + 三因子组合扫描
果 = 该温区下 fwd_ret60 前 20% 分位（相对，避免绝对阈值稀疏）
"""
import json
import sqlite3
import numpy as np
import pandas as pd
from datetime import datetime
from itertools import combinations

DATA_DIR = "/home/hmo/MoFin/data"
OUT_JSON = f"{DATA_DIR}/b_group_candidates.json"


def load_regime_map(market="a"):
    conn = sqlite3.connect("/home/hmo/MoFin/data/mofin.db", timeout=10)
    rows = conn.execute("SELECT date, regime FROM market_regime WHERE market=?", (market,)).fetchall()
    conn.close()
    return {d: r for d, r in rows}


def load_panel(market):
    path = "/tmp/panel_12d_hk.pkl" if market == "hk" else "/tmp/panel_12d.pkl"
    p = pd.read_pickle(path)
    p = p.sort_values(["code", "date"]).reset_index(drop=True)
    p["fwd_ret60"] = p.groupby("code")["close"].transform(lambda x: x.shift(-60) / x - 1) * 100
    return p


def scan3(market, regime, panel, min_n=500):
    """三因子组合扫描：相对分位果"""
    rm = load_regime_map(market)
    p = panel.copy()
    p["_regime"] = p["date"].map(rm)
    sub = p[p["_regime"] == regime].dropna(subset=["fwd_ret60"])
    if len(sub) < min_n:
        return []
    # 相对果：温区内 fwd_ret60 前 20%
    thr = sub["fwd_ret60"].quantile(0.80)
    sub["is_good"] = (sub["fwd_ret60"] >= thr).astype(int)
    br = 20.0  # 相对分位定义，基线恒 20%
    print(f"[{market}/{regime}] 样本{len(sub)} 果阈值60日+{thr:.0f}%")

    # 因子池（方向：小市值/低估值/超跌/放量/企稳/低动量）
    factor_defs = {
        "mcap_q": ("<", 0.5), "pe_q": ("<", 0.5), "pb_q": ("<", 0.5),
        "bias60": ("<", -5), "rsi": ("<", 50), "dist_lo20": (">", 3),
        "vol_ratio": (">", 1.0), "mkt_ret20": ("<", 0), "ret20": ("<", 0),
        "sec_ret20": ("<", 0), "flow5": (">", 0), "news3": (">=", 1),
        "ret5": (">", -3), "mkt_rsi": ("<", 50),
    }
    # 单条件
    single = []
    for feat, (op, val) in factor_defs.items():
        if feat not in sub.columns:
            continue
        cond = sub[feat] < val if op == "<" else sub[feat] > val
        m = sub[cond]
        if len(m) < 200:
            continue
        rate = m["is_good"].mean() * 100
        if rate > 23:  # 相对基线20% +3pp
            single.append((feat, round(rate, 1), len(m), round(rate - 20, 1)))
    single.sort(key=lambda x: -x[3])
    print("  单条件:", single[:4])

    # 三因子组合（从单条件超额>2pp 里取 6 个，C(6,3)=20 组合）
    pool = [s[0] for s in single if s[3] > 2][:6]
    results = []
    for combo in combinations(pool, 3):
        cond = pd.Series(True, index=sub.index)
        for feat in combo:
            op, val = factor_defs[feat]
            cond &= (sub[feat] < val) if op == "<" else (sub[feat] > val)
        m = sub[cond]
        if len(m) < 200:
            continue
        rate = m["is_good"].mean() * 100
        avg = m["fwd_ret60"].mean()
        results.append(({f: factor_defs[f] for f in combo}, len(m), round(rate, 1),
                        round(avg, 1), round(rate - 20, 1)))
    results.sort(key=lambda x: -x[4])
    return results[:5]


def to_entry(cond_dict):
    entry = {}
    for feat, (op, val) in cond_dict.items():
        key = feat + ("_min" if op == ">" else "_max")
        entry[key] = float(val)
    return entry


def mine(market="a", regimes=None):
    regimes = regimes or ["trend_up", "choppy", "trend_down"]
    panel = load_panel(market)
    out = {"market": market, "mined_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "candidates": []}
    for rg in regimes:
        combos = scan3(market, rg, panel)
        for cond, n, rate, avg, extra in combos[:3]:
            cand = {
                "regime": rg, "market": market, "group": "B", "status": "candidate",
                "entry": to_entry(cond), "trades_est": n, "good_rate": rate,
                "avg60": avg, "excess_pp": extra,
                "hypothesis": f"[{rg}] 由果及因三因子: {list(cond.keys())} → 60日前20%占比{rate}%(超额+{extra}pp)",
            }
            out["candidates"].append(cand)
            print(f"  [{rg}] {list(cond.keys())} n={n} 好果率{rate}% 超额+{extra}pp")
    return out


if __name__ == "__main__":
    import sys
    market = sys.argv[1] if len(sys.argv) > 1 else "a"
    res = mine(market)
    with open(OUT_JSON, "w", encoding="utf-8") as f:
        json.dump(res, f, ensure_ascii=False, indent=1)
    print(f"写入 {OUT_JSON}: {len(res['candidates'])} 个候选")