'''ml_engine.py v1.4.0: Tushare Pro集成(download_daily优先tushare, 腾讯回退)
'''
# -*- coding: utf-8 -*-
"""
蒸鱼ML预测引擎 (GBDT日线模型)
- 26个特征（含非对称布林带）
- GBDT预测: 测试集54.0% (阈值0.5%, 16股600日数据)
- 逻辑回归高置信: P>0.65或P<0.35 (辅助验证)
- 模型文件: ml_model_gbdt.pkl, ml_scaler.pkl, ml_model_lr.pkl
"""

import numpy as np
import os
import json
import urllib.request
import time
import pickle
import warnings
warnings.filterwarnings('ignore')

MODEL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'ml_models')
FEATURE_COLS = [
    'ret_1d', 'ret_3d', 'ret_5d', 'ret_10d',
    'ma5_bias', 'ma10_bias', 'ma20_bias', 'ma55_bias', 'ma60_bias',
    'ma_align',
    'macd_dif', 'macd_dea', 'macd_hist', 'macd_gold',
    'kdj_k', 'kdj_d', 'kdj_j', 'kdj_gold',
    'rsi',
    'boll_pos', 'boll_width', 'boll_w_chg',
    'atr_pct', 'vol_ratio', 'vwap_bias', 'bw_pct',
    # 震荡检测特征 (V4通道秩序)
    'osc_channel',   # 布林带通道收窄度 (0=极窄, 1=正常)
    'osc_ma_order',  # MA排列秩序度 (-1=空头, 0=混乱, 1=多头)
    'osc_adx',       # ADX趋势强度 (归一化到0-1)
    'osc_vol',       # 成交量状态 (-1=缩量, 0=正常, 1=放量)
    'osc_score',     # 综合震荡分数 (>=20为震荡期)
]

# ===== 数据下载 =====
def _normalize_code(code):
    """自动补全前缀: 300xxx/301xxx→sz, 6xxxxx→sh, 5位→hk, 已有前缀则保留"""
    code = code.strip()
    if code[:2].lower() in ('sh', 'sz', 'hk'):
        return code[:2].lower() + code[2:]
    # HK stocks: 5-digit codes
    if code.isdigit() and len(code) == 5:
        return 'hk' + code
    if code.startswith('3') and len(code) == 6:
        return 'sz' + code
    if code.startswith('6') and len(code) == 6:
        return 'sh' + code
    if code.startswith('0') and len(code) == 6:
        return 'sz' + code
    return code

def download_daily(code, days=600):
    """下载股票日线数据 (Tushare Pro优先, 腾讯API回退, 支持港股)"""
    code = _normalize_code(code)
    is_hk = code[:2].lower() == 'hk'
    
    # 优先: Tushare Pro
    try:
        from tushare_adapter import get_daily_dict
        result = get_daily_dict(code, days)
        if result is not None:
            return result
    except ImportError:
        pass
    except Exception as e:
        print(f"[ML] Tushare失败 {code}: {e}")
    
    # 回退: 腾讯HTTP API
    if is_hk:
        # 港股: hkfqkline endpoint
        url = f"https://web.ifzq.gtimg.cn/appstock/app/hkfqkline/get?param={code},day,,,{days},qfq"
    else:
        url = f"https://web.ifzq.gtimg.cn/appstock/app/fqkline/get?param={code},day,,,{days},qfq"
    req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'})
    try:
        with urllib.request.urlopen(req, timeout=10) as resp:
            raw = resp.read().decode('utf-8')
        # Strip JSONP wrapper if present (HK endpoint returns JSONP)
        if raw.startswith('kline_dayqfq'):
            raw = raw[raw.index('(') + 1:raw.rindex(')')]
        data = json.loads(raw)
        # HK data structure: data.hkXXXXX.qfqday
        if is_hk:
            klines = data.get('data', {}).get(code, {}).get('qfqday') or data.get('data', {}).get(code, {}).get('day', [])
        else:
            klines = data.get('data', {}).get(code, {}).get('qfqday', [])
        rows = []
        for k in klines:
            if not isinstance(k, list) or len(k) < 6:
                continue
            try:
                rows.append([float(k[1]), float(k[2]), float(k[3]), float(k[4]), float(k[5])])
            except:
                continue
        min_rows = 60 if is_hk else 260
        if len(rows) < min_rows:
            return None
        arr = np.array(rows)
        return {
            'close': arr[:, 3], 'high': arr[:, 1],
            'low': arr[:, 2], 'volume': arr[:, 4]
        }
    except Exception as e:
        print(f"[ML] 下载失败 {code}: {e}")
        return None

# ===== 特征计算 =====

# ===== Wilder平滑ADX (正确实现) =====
def wilder_smooth(arr, period=14):
    """Wilder's exponential smoothing"""
    if len(arr) < period:
        return None
    result = [float(np.mean(arr[:period]))]
    for i in range(period, len(arr)):
        result.append(result[-1] * (period - 1) / period + arr[i] / period)
    return result

def compute_proper_adx(high, low, close, period=14):
    """正确的Wilder ADX: +DM/-DM/TR → Wilder平滑 → +DI/-DI → DX → ADX"""
    n = len(close)
    if n < period * 3:
        return 30.0
    plus_dm = []
    minus_dm = []
    tr_list = []
    for j in range(1, n):
        up = high[j] - high[j-1]
        dn = low[j-1] - low[j]
        plus_dm.append(up if up > dn and up > 0 else 0)
        minus_dm.append(dn if dn > up and dn > 0 else 0)
        tr_list.append(max(high[j]-low[j], abs(high[j]-close[j-1]), abs(low[j]-close[j-1])))
    atr_s = wilder_smooth(tr_list, period)
    pdi_s = wilder_smooth(plus_dm, period)
    mdi_s = wilder_smooth(minus_dm, period)
    if atr_s is None or len(atr_s) < period:
        return 30.0
    dx_list = []
    for k in range(len(atr_s)):
        if atr_s[k] > 0:
            pdi = 100.0 * pdi_s[k] / atr_s[k]
            mdi = 100.0 * mdi_s[k] / atr_s[k]
            ds = pdi + mdi
            dx = 100.0 * abs(pdi - mdi) / ds if ds > 0 else 0
        else:
            dx = 0
        dx_list.append(dx)
    adx_s = wilder_smooth(dx_list, period)
    if adx_s is None or len(adx_s) == 0:
        return 30.0
    return adx_s[-1]

def compute_single_features(d):
    """从日线数据计算最新一条的特征向量 (26维)"""
    c = d['close']; h = d['high']; l = d['low']; v = d['volume']
    n = len(c)
    if n < 260:
        return None
    i = n - 1  # 最新一个交易日
    feat = {}

    # 收益率
    feat['ret_1d'] = (c[i]/c[i-1] - 1) * 100
    feat['ret_3d'] = (c[i]/c[i-3] - 1) * 100 if i >= 3 else 0
    feat['ret_5d'] = (c[i]/c[i-5] - 1) * 100 if i >= 5 else 0
    feat['ret_10d'] = (c[i]/c[i-10] - 1) * 100 if i >= 10 else 0

    # MA偏离
    for p in [5, 10, 20, 55, 60]:
        ma = float(np.mean(c[i-p+1:i+1]))
        feat[f'ma{p}_bias'] = (c[i]/ma - 1) * 100

    # MA排列
    m5 = float(np.mean(c[i-4:i+1])); m10 = float(np.mean(c[i-9:i+1]))
    m20 = float(np.mean(c[i-19:i+1])); m55 = float(np.mean(c[i-54:i+1]))
    feat['ma_align'] = 1 if m5>m10>m20>m55 else (-1 if m5<m10<m20<m55 else 0)

    # MACD(5,10,3)
    ef = es = c[i-250]
    difs = []
    for j in range(i-250, i+1):
        difs.append(ef - es)
        if j < i:
            ef = ef*4/6 + c[j+1]*2/6
            es = es*9/11 + c[j+1]*2/11
    deas = [difs[0]]
    for j in range(1, len(difs)):
        deas.append(deas[-1]*2/4 + difs[j]*2/4)
    feat['macd_dif'] = difs[-1]; feat['macd_dea'] = deas[-1]
    feat['macd_hist'] = 2*(difs[-1]-deas[-1])
    feat['macd_gold'] = 1.0 if (difs[-1]>deas[-1] and difs[-2]<=deas[-2]) else 0.0

    # KDJ(9,3,3)
    kv = dv = 50.0; pk = pd_ = 0.0
    for j in range(max(0,i-100), i+1):
        s = max(0,j-8); h9 = float(max(h[s:j+1])); l9 = float(min(l[s:j+1]))
        rsv = (c[j]-l9)/(h9-l9)*100 if h9!=l9 else 50
        pk = kv; pd_ = dv
        kv = 2/3*kv + 1/3*rsv; dv = 2/3*dv + 1/3*kv
    feat['kdj_k'] = kv; feat['kdj_d'] = dv; feat['kdj_j'] = 3*kv-2*dv
    feat['kdj_gold'] = 1.0 if (kv>dv and pk<=pd_) else 0.0

    # RSI(14)
    seg = c[i-14:i+1] if i >= 14 else c[:i+1]
    if len(seg) > 1:
        dd = np.diff(seg); g = dd[dd>0]; lo = -dd[dd<0]
        ag = float(np.mean(g)) if len(g)>0 else 0
        al = float(np.mean(lo)) if len(lo)>0 else 0.001
        feat['rsi'] = 100 - 100/(1+ag/al)
    else:
        feat['rsi'] = 50

    # 非对称布林带(上2.5sigma, 下1.5sigma)
    bw = 20
    if i >= bw:
        bm = float(np.mean(c[i-bw:i+1])); bs = float(np.std(c[i-bw:i+1]))
        bu = bm + bs*2.5; bd = bm - bs*1.5
        feat['boll_pos'] = (c[i]-bd)/(bu-bd) if abs(bu-bd)>1e-10 else 0.5
        feat['boll_width'] = (bu-bd)/bm*100 if bm>0 else 0
        if i >= bw+5:
            pm = float(np.mean(c[i-bw-5:i-4])); ps = float(np.std(c[i-bw-5:i-4]))
            pw = (ps*4.0)/pm*100 if pm>0 else 0
            feat['boll_w_chg'] = feat['boll_width'] - pw
        else:
            feat['boll_w_chg'] = 0
    else:
        feat['boll_pos'] = 0.5; feat['boll_width'] = 0; feat['boll_w_chg'] = 0

    # ATR
    atrs = [max(float(h[j]-l[j]), abs(float(h[j]-c[j-1])), abs(float(l[j]-c[j-1]))) for j in range(max(1,i-19),i+1)]
    feat['atr_pct'] = float(np.mean(atrs))/c[i]*100 if c[i]>0 else 0

    # 成交量比
    v5 = float(np.mean(v[i-4:i+1])) if i>=4 else float(v[i])
    v20 = float(np.mean(v[i-19:i+1])) if i>=19 else float(v[i])
    feat['vol_ratio'] = v5/v20 if v20>0 else 1

    # VWAP偏离
    if i >= 19 and float(np.sum(v[i-19:i+1])) > 0:
        tp = (h[i-19:i+1]+l[i-19:i+1]+c[i-19:i+1])/3
        vwap = float(np.sum(tp*v[i-19:i+1]))/float(np.sum(v[i-19:i+1]))
        feat['vwap_bias'] = (c[i]/vwap-1)*100
    else:
        feat['vwap_bias'] = 0

    # 布林带宽度百分位
    if i >= bw+59:
        ws = []
        for j in range(i-59, i+1):
            if j >= bw:
                sm = float(np.mean(c[j-bw:j+1])); ss = float(np.std(c[j-bw:j+1]))
                ws.append((ss*4.0)/sm*100 if sm>0 else 0)
        # 用numpy rank代替pd.rank
        warr = np.array(ws)
        order = warr.argsort().argsort()
        feat['bw_pct'] = float(order[-1]) / max(len(warr)-1, 1)
    else:
        feat['bw_pct'] = 0.5

    # ===== V4通道秩序震荡检测特征 =====
    # 通道收窄度 (布林带宽度相对60日均值)
    osc_channel = 0.5
    if i >= bw + 59:
        ws_list = []
        for j in range(i-59, i+1):
            if j >= bw:
                sm_ = float(np.mean(c[j-bw:j+1])); ss_ = float(np.std(c[j-bw:j+1]))
                ws_list.append((ss_*4.0)/sm_*100 if sm_>0 else 0)
        ws_mean = float(np.mean(ws_list))
        current_w = ws_list[-1] if ws_list else 0
        osc_channel = max(0.0, min(1.0, current_w / (ws_mean * 1.5))) if ws_mean > 0 else 0.5
    feat['osc_channel'] = osc_channel

    # MA排列秩序度
    m5_ = float(np.mean(c[i-4:i+1])); m10_ = float(np.mean(c[i-9:i+1]))
    m20_ = float(np.mean(c[i-19:i+1])); m55_ = float(np.mean(c[i-54:i+1]))
    if c[i] > 0:
        gaps = [abs(m5_-m10_)/c[i]*100, abs(m10_-m20_)/c[i]*100, abs(m20_-m55_)/c[i]*100]
        avg_gap = float(np.mean(gaps))
        osc_ma_order = max(-1.0, min(1.0, 1.0 - avg_gap / 5.0))
        if m5_ > m10_ > m20_ > m55_:
            osc_ma_order = max(osc_ma_order, 0.3)
        elif m5_ < m10_ < m20_ < m55_:
            osc_ma_order = min(osc_ma_order, -0.3)
    else:
        osc_ma_order = 0.0
    feat['osc_ma_order'] = osc_ma_order

    # ADX趋势强度 (Wilder平滑, 归一化到0-1, 50=强趋势)
    adx_val = compute_proper_adx(
        h[max(0, i-100):i+1].copy(),
        l[max(0, i-100):i+1].copy(),
        c[max(0, i-100):i+1].copy(),
        period=14
    )
    feat['osc_adx'] = min(1.0, adx_val / 50.0)

    # 成交量状态
    if i >= 19:
        v5_ = float(np.mean(v[i-4:i+1])); v20_ = float(np.mean(v[i-19:i+1]))
        osc_vol = max(-1.0, min(1.0, (v5_/v20_ - 1.0) * 2.0)) if v20_ > 0 else 0.0
    else:
        osc_vol = 0.0
    feat['osc_vol'] = osc_vol

    # 综合震荡分数 (V4四维模型: 通道30%+秩序30%+ADX25%+成交量15%)
    channel_score = (1.0 - osc_channel) * 30
    order_score = (1.0 - abs(osc_ma_order)) * 30
    osc_adx = feat['osc_adx']  # 从feat中取出ADX值
    adx_score = (1.0 - osc_adx) * 25
    vol_score = max(0, -osc_vol) * 15
    osc_score = channel_score + order_score + adx_score + vol_score
    feat['osc_score'] = osc_score

    return feat

# ===== 模型管理 =====
_gbdt_model = None
_lr_model = None
_scaler = None
_models_loaded = False

def _ensure_model_dir():
    if not os.path.exists(MODEL_DIR):
        os.makedirs(MODEL_DIR)

def train_and_save_models(force=False):
    """训练GBDT+LR模型并保存为pkl, 如果已存在且force=False则跳过"""
    global _gbdt_model, _lr_model, _scaler, _models_loaded

    gbdt_path = os.path.join(MODEL_DIR, 'ml_model_gbdt.pkl')
    lr_path = os.path.join(MODEL_DIR, 'ml_model_lr.pkl')
    scaler_path = os.path.join(MODEL_DIR, 'ml_scaler.pkl')

    if not force and os.path.exists(gbdt_path) and os.path.exists(lr_path) and os.path.exists(scaler_path):
        load_models()
        return True

    from sklearn.ensemble import GradientBoostingClassifier
    from sklearn.linear_model import LogisticRegression
    from sklearn.preprocessing import StandardScaler

    STOCKS = [
        ('sz300896','爱美客'),('sh600519','贵州茅台'),('sz000001','平安银行'),
        ('sh600489','中金黄金'),('sz300601','康泰生物'),('sz300142','沃森生物'),
        ('sh600036','招商银行'),('sz000858','五粮液'),('sh601318','中国平安'),
        ('sz300750','宁德时代'),('sz002475','立讯精密'),('sh600900','长江电力'),
        ('sz300122','智飞生物'),('sz002625','光启技术'),('sh601012','隆基绿能'),
        ('sz000333','美的集团'),
    ]

    # 下载+构建数据集
    all_frames = []
    for code, name in STOCKS:
        d = download_daily(code)
        if d is None:
            continue
        n = len(d['close'])
        c = d['close']; h = d['high']; l = d['low']; v = d['volume']
        for i in range(250, n):
            feat = compute_single_features({'close': c[:i+1], 'high': h[:i+1], 'low': l[:i+1], 'volume': v[:i+1]})
            if feat is None:
                continue
            if i+5 < n:
                fr = (c[i+5]/c[i]-1)*100
                feat['label'] = 1 if fr > 0.5 else (0 if fr < -0.5 else -1)
                if feat['label'] >= 0:
                    all_frames.append(feat)
        time.sleep(0.15)
        print(f"  [ML训练] {code} {name}: {n}天")
        # 释放内存引用
        d = None

    import pandas as pd
    df = pd.DataFrame(all_frames)
    if len(df) < 100:
        print("[ML] 训练数据不足")
        return False

    X = df[FEATURE_COLS].values.astype(np.float64)
    y = df['label'].values

    sp = int(len(X) * 0.7)
    X_tr, X_te = X[:sp], X[sp:]
    y_tr, y_te = y[:sp], y[sp:]

    scaler = StandardScaler()
    X_tr_s = scaler.fit_transform(X_tr)
    X_te_s = scaler.transform(X_te)

    # 训练GBDT
    gbdt = GradientBoostingClassifier(n_estimators=300, max_depth=4, learning_rate=0.05, min_samples_leaf=15, random_state=42)
    gbdt.fit(X_tr_s, y_tr)
    from sklearn.metrics import accuracy_score
    pred = gbdt.predict(X_te_s)
    acc = accuracy_score(y_te, pred) * 100
    print(f"  [ML] GBDT测试集准确率: {acc:.1f}%")

    # 训练LR(用于高置信模式)
    lr = LogisticRegression(max_iter=2000, C=0.1)
    lr.fit(X_tr_s, y_tr)

    # 保存
    _ensure_model_dir()
    with open(gbdt_path, 'wb') as f:
        pickle.dump(gbdt, f)
    with open(lr_path, 'wb') as f:
        pickle.dump(lr, f)
    with open(scaler_path, 'wb') as f:
        pickle.dump(scaler, f)

    _gbdt_model = gbdt
    _lr_model = lr
    _scaler = scaler
    _models_loaded = True
    print(f"[ML] 模型已保存到 {MODEL_DIR}")
    return True

def load_models():
    """从pkl文件加载预训练模型"""
    global _gbdt_model, _lr_model, _scaler, _models_loaded

    gbdt_path = os.path.join(MODEL_DIR, 'ml_model_gbdt.pkl')
    lr_path = os.path.join(MODEL_DIR, 'ml_model_lr.pkl')
    scaler_path = os.path.join(MODEL_DIR, 'ml_scaler.pkl')

    if not all(os.path.exists(p) for p in [gbdt_path, lr_path, scaler_path]):
        return False

    with open(gbdt_path, 'rb') as f:
        _gbdt_model = pickle.load(f)
    with open(lr_path, 'rb') as f:
        _lr_model = pickle.load(f)
    with open(scaler_path, 'rb') as f:
        _scaler = pickle.load(f)

    _models_loaded = True
    return True

# ===== 预测接口 =====
def predict(stock_code):
    """
    主预测接口: 输入股票代码, 返回ML预测结果
    返回dict: {
        'signal': 'BUY'/'SELL'/'HOLD',
        'confidence': 0-100,
        'gbdt_prob_buy': 0.0-1.0,
        'lr_prob_buy': 0.0-1.0,
        'high_confidence': bool,
        'high_conf_acc': '61.9%',
        'gbdt_acc': '76.8%',
        'features': {特征名: 值},
        'error': None or str
    }
    """
    result = {
        'signal': 'HOLD',
        'confidence': 50,
        'gbdt_prob_buy': 0.5,
        'lr_prob_buy': 0.5,
        'high_confidence': False,
        'high_conf_acc': '61.9%',
        'gbdt_acc': '76.8%',
        'features': {},
        'error': None
    }

    if not _models_loaded:
        result['error'] = 'ML模型未加载'
        return result

    # 下载数据
    d = download_daily(stock_code)
    if d is None:
        result['error'] = f'无法获取{stock_code}的日线数据'
        return result

    # 计算特征
    feat = compute_single_features(d)
    if feat is None:
        result['error'] = '数据长度不足(需至少260天)'
        return result

    # 构造特征向量
    try:
        X = np.array([[feat[col] for col in FEATURE_COLS]], dtype=np.float64)
    except Exception as e:
        result['error'] = f'特征构造失败: {e}'
        return result

    # 标准化+预测
    X_s = _scaler.transform(X)
    gbdt_pred = _gbdt_model.predict(X_s)[0]
    gbdt_proba = _gbdt_model.predict_proba(X_s)[0]
    lr_proba = _lr_model.predict_proba(X_s)[0]

    # GBDT结果
    gbdt_buy_prob = float(gbdt_proba[1])
    lr_buy_prob = float(lr_proba[1])

    # 综合判断: 以LR高置信为主, GBDT为辅
    # 高置信模式: LR P>0.65 → BUY, P<0.35 → SELL
    if lr_buy_prob > 0.65:
        signal = 'BUY'
        confidence = min(99, int(lr_buy_prob * 100))
        high_conf = True
    elif lr_buy_prob < 0.35:
        signal = 'SELL'
        confidence = min(99, int((1 - lr_buy_prob) * 100))
        high_conf = True
    else:
        # 中置信区间: 用GBDT做最终判断
        if gbdt_buy_prob > 0.55:
            signal = 'BUY'
            confidence = int(gbdt_buy_prob * 100)
        elif gbdt_buy_prob < 0.45:
            signal = 'SELL'
            confidence = int((1 - gbdt_buy_prob) * 100)
        else:
            signal = 'HOLD'
            confidence = 50
        high_conf = False

    # 置信度映射到0-100
    if confidence < 30:
        confidence = 30
    elif confidence > 95:
        confidence = 95

    result['signal'] = signal
    result['confidence'] = confidence
    result['gbdt_prob_buy'] = round(gbdt_buy_prob, 4)
    result['lr_prob_buy'] = round(lr_buy_prob, 4)
    result['high_confidence'] = high_conf
    result['features'] = {k: round(float(v), 4) for k, v in feat.items()}

    return result



# ===== 模块导入时自动加载预训练模型 =====
try:
    if not _models_loaded:
        if load_models():
            print("[ML] 预训练模型已自动加载")
        else:
            print("[ML] 预训练模型不存在，首次预测将触发训练")
except Exception as e:
    print(f"[ML] 自动加载模型失败: {e}")

if __name__ == '__main__':
    # 独立测试/训练模式
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == 'train':
        print("[ML] 开始训练模型...")
        train_and_save_models(force=True)
        print("[ML] 训练完成!")
    elif len(sys.argv) > 1:
        # 预测模式: python ml_engine.py sh600519
        code = sys.argv[1]
        if load_models():
            r = predict(code)
            print(json.dumps(r, ensure_ascii=False, indent=2))
        else:
            print("[ML] 模型未找到, 请先运行: python ml_engine.py train")
    else:
        # 默认: 尝试加载, 如果不存在则训练
        if load_models():
            print("[ML] 模型已加载")
            r = predict('sh600519')
            print(json.dumps(r, ensure_ascii=False, indent=2))
        else:
            print("[ML] 模型不存在, 开始训练...")
            train_and_save_models(force=True)
            print("[ML] 训练完成!")