#!/usr/bin/env python3
"""MoFin 回测框架 v2 — 从 stock_daily 原始行情计算全部指标，模拟全流程"""
import sqlite3, json, math
from datetime import datetime, timedelta
from collections import defaultdict

DB_PATH = "/home/hmo/MoFin/data/mofin.db"

# ── 技术指标计算（纯 Python，不依赖第三方库）──

def calc_ma(series, n):
    """简单移动平均"""
    result = []
    for i in range(len(series)):
        if i < n - 1:
            result.append(None)
        else:
            result.append(sum(series[i-n+1:i+1]) / n)
    return result

def calc_ema(series, n):
    """指数移动平均"""
    result = []
    k = 2 / (n + 1)
    for i in range(len(series)):
        if i == 0:
            result.append(series[i])
        else:
            result.append(series[i] * k + result[-1] * (1 - k))
    return result

def calc_rsi(series, n=14):
    """RSI"""
    deltas = [series[i] - series[i-1] for i in range(1, len(series))]
    gains = [d if d > 0 else 0 for d in deltas]
    losses = [-d if d < 0 else 0 for d in deltas]
    result = [None] * (n + 1)
    avg_gain = sum(gains[:n]) / n
    avg_loss = sum(losses[:n]) / n
    if avg_loss == 0:
        result.append(100)
    else:
        rs = avg_gain / avg_loss
        result.append(100 - 100 / (1 + rs))
    for i in range(n, len(gains)):
        avg_gain = (avg_gain * (n-1) + gains[i]) / n
        avg_loss = (avg_loss * (n-1) + losses[i]) / n
        if avg_loss == 0:
            result.append(100)
        else:
            rs = avg_gain / avg_loss
            result.append(100 - 100 / (1 + rs))
    # Pad front to match original length
    while len(result) < len(series):
        result.insert(0, None)
    return result[:len(series)]

def calc_macd(series, fast=12, slow=26, signal=9):
    """MACD: 返回 macd_line, signal_line, histogram"""
    ema_fast = calc_ema(series, fast)
    ema_slow = calc_ema(series, slow)
    macd_line = [ema_fast[i] - ema_slow[i] if ema_fast[i] is not None and ema_slow[i] is not None else None for i in range(len(series))]
    # Recompute signal line from macd_line (only non-None values)
    macd_clean = [v for v in macd_line if v is not None]
    if not macd_clean:
        return [None]*len(series), [None]*len(series), [None]*len(series)
    signal_raw = calc_ema(macd_clean, signal)
    signal_line = [None] * (len(series) - len(macd_clean)) + signal_raw
    histogram = [macd_line[i] - signal_line[i] if macd_line[i] is not None and signal_line[i] is not None else None for i in range(len(series))]
    return macd_line, signal_line, histogram

def calc_tr(highs, lows, closes):
    """True Range"""
    tr = [highs[0] - lows[0]]
    for i in range(1, len(highs)):
        hl = highs[i] - lows[i]
        hc = abs(highs[i] - closes[i-1])
        lc = abs(lows[i] - closes[i-1])
        tr.append(max(hl, hc, lc))
    return tr

def calc_atr(highs, lows, closes, n=14):
    """ATR"""
    tr = calc_tr(highs, lows, closes)
    return calc_ema(tr, n)

def calc_trend_strength(highs, lows, closes, n=14):
    """简易趋势强度（替代 ADX）: (DI+ - DI-) / (DI+ + DI-) 归一化"""
    tr = calc_tr(highs, lows, closes)
    up = [highs[i] - highs[i-1] for i in range(1, len(highs))]
    down = [lows[i-1] - lows[i] for i in range(1, len(lows))]
    # Directional indicators (raw)
    di_plus_raw = [0.0] * len(up)
    di_minus_raw = [0.0] * len(up)
    for i in range(len(up)):
        if up[i] > down[i] and up[i] > 0:
            di_plus_raw[i] = up[i]
        if down[i] > up[i] and down[i] > 0:
            di_minus_raw[i] = down[i]
        # Normalize by TR
        tr_val = tr[i+1] if (i+1) < len(tr) else tr[-1]
        if tr_val > 0:
            di_plus_raw[i] = di_plus_raw[i] / tr_val * 100
            di_minus_raw[i] = di_minus_raw[i] / tr_val * 100
    
    # SMA of DMI
    result = []
    for i in range(len(di_plus_raw)):
        if i < n - 1:
            result.append(None)
        else:
            avg_plus = sum(di_plus_raw[i-n+1:i+1]) / n
            avg_minus = sum(di_minus_raw[i-n+1:i+1]) / n
            if avg_plus + avg_minus > 0:
                dx = abs(avg_plus - avg_minus) / (avg_plus + avg_minus) * 100
            else:
                dx = 0
            result.append(dx)
    # Pad front to match original highs length
    return [None] * (len(highs) - len(result)) + result

def calc_obv(closes, volumes):
    """OBV"""
    obv = [volumes[0]]
    for i in range(1, len(closes)):
        if closes[i] > closes[i-1]:
            obv.append(obv[-1] + volumes[i])
        elif closes[i] < closes[i-1]:
            obv.append(obv[-1] - volumes[i])
        else:
            obv.append(obv[-1])
    return obv

def calc_roc(series, n=10):
    """Rate of Change"""
    result = []
    for i in range(len(series)):
        if i < n:
            result.append(None)
        else:
            result.append((series[i] - series[i-n]) / series[i-n] * 100 if series[i-n] != 0 else 0)
    return result

def prepare_bars(code, start_date, end_date):
    """获取原始数据并计算全部指标"""
    conn = sqlite3.connect(DB_PATH)
    rows = conn.execute("""
        SELECT date, open, close, high, low, volume, amount
        FROM stock_daily WHERE code=? AND date>=? AND date<=?
        ORDER BY date
    """, (code, start_date, end_date)).fetchall()
    conn.close()
    if not rows or len(rows) < 30:
        return None
    dates = [r[0] for r in rows]
    opens = [r[1] for r in rows]
    closes = [r[2] for r in rows]
    highs = [r[3] for r in rows]
    lows = [r[4] for r in rows]
    volumes = [r[5] for r in rows]
    amounts = [r[6] if len(r) > 6 else None for r in rows]
    
    # 计算全部指标
    ma5 = calc_ma(closes, 5)
    ma10 = calc_ma(closes, 10)
    ma20 = calc_ma(closes, 20)
    ma60 = calc_ma(closes, 60)
    rsi = calc_rsi(closes)
    macd_line, signal_line, macd_hist = calc_macd(closes)
    atr = calc_atr(highs, lows, closes)
    trend_strength = calc_trend_strength(highs, lows, closes)
    obv = calc_obv(closes, volumes)
    roc = calc_roc(closes)

    bars = []
    for i in range(len(dates)):
        # ── 枢轴S/R（与实盘 technical_analysis.calc_support_resistance 同算法）──
        piv = {}
        if i >= 1:
            h, l, c = highs[i], lows[i], closes[i]
            yc = closes[i-1]
            # 多日波幅（近20日）
            win = bars[max(0, i-19):i] if i > 0 else []
            multi_high = max([h] + [x.get('high', 0) or 0 for x in win])
            multi_low = min([l] + [x.get('low', 1e9) or 1e9 for x in win])
            eff_range = max(h - l, multi_high - multi_low, c * 0.05)
            tp = (c - multi_low) / (multi_high - multi_low) if multi_high > multi_low else 0.5
            if tp > 0.8 or tp < 0.2:
                eff_range = max(eff_range, c * 0.08)
            pp = (h + l + c) / 3
            s1 = 2 * pp - h
            s2 = pp - eff_range
            r1 = 2 * pp - l
            r2 = pp + eff_range
            if yc < s1: s1 = yc
            if yc > r1: r1 = yc
            piv = {'pivot': pp, 'weak_support': s1, 'strong_support': s2,
                   'weak_resist': r1, 'strong_resist': r2}

        bars.append({
            'date': dates[i],
            'open': opens[i],
            'close': closes[i],
            'high': highs[i],
            'low': lows[i],
            'volume': volumes[i],
            'amount': amounts[i] if i < len(amounts) else None,
            'ma5': ma5[i] if i < len(ma5) else None,
            'ma10': ma10[i] if i < len(ma10) else None,
            'ma20': ma20[i] if i < len(ma20) else None,
            'ma60': ma60[i] if i < len(ma60) else None,
            'rsi': rsi[i] if i < len(rsi) else None,
            'macd': macd_line[i] if i < len(macd_line) else None,
            'macd_signal': signal_line[i] if i < len(signal_line) else None,
            'macd_hist': macd_hist[i] if i < len(macd_hist) else None,
            'atr': atr[i] if i < len(atr) else None,
            'adx': trend_strength[i] if i < len(trend_strength) else None,
            'obv': obv[i] if i < len(obv) else None,
            'roc': roc[i] if i < len(roc) else None,
            **piv,
        })
    return bars


def compute_single_score(bars):
    """对单个股票计算多因子评分（与 MoFin 现有一致）"""
    if not bars or len(bars) < 15:
        return None
    last = bars[-1]
    prev5 = bars[-5] if len(bars) >= 5 else bars[-2] if len(bars) >= 2 else last
    prev10 = bars[-10] if len(bars) >= 10 else bars[0]
    
    def v(val): return val if val is not None else 0
    score = {'trend': 0, 'momentum': 0, 'volume': 0, 'volatility': 0, 'risk': 0}
    
    # 趋势分 (0-25)
    ts = 0
    adx = v(last.get('adx'))
    if adx > 25: ts += 10
    elif adx > 20: ts += 5
    mh = v(last.get('macd_hist'))
    pmh = v(prev5.get('macd_hist'))
    if mh > 0 and mh > pmh: ts += 10
    elif mh > 0: ts += 5
    elif mh < 0 and mh < pmh: ts -= 5
    ma5 = v(last.get('ma5'))
    ma10 = v(last.get('ma10'))
    ma20 = v(last.get('ma20'))
    if ma5 > ma10 > ma20 and ma5 > 0 and ma10 > 0: ts += 5
    elif ma5 < ma10 < ma20 and ma5 > 0 and ma10 > 0: ts -= 5
    score['trend'] = max(0, min(25, ts))
    
    # 动量分 (0-20)
    ms = 0
    roc = v(last.get('roc'))
    if roc > 5: ms += 5
    elif roc > 2: ms += 2
    elif roc < -5: ms -= 5
    elif roc < -2: ms -= 2
    rsi = v(last.get('rsi'))
    if 30 <= rsi <= 70: ms += 5
    if 50 < rsi <= 65: ms += 3
    elif rsi > 70: ms -= 3
    elif rsi < 30: ms -= 3
    close = v(last.get('close'))
    if ma20 > 0:
        pct = (close - ma20) / ma20 * 100
        if -3 <= pct <= 5: ms += 7
        elif pct < -10: ms -= 3
    score['momentum'] = max(0, min(20, ms))
    
    # 量能分 (0-20)
    vs = 0
    obv_cur = v(last.get('obv'))
    obv_prev = v(prev5.get('obv'))
    vol_cur = v(last.get('volume'))
    vol_prev = v(prev5.get('volume'))
    if obv_cur > obv_prev and vol_cur > 0 and vol_prev > 0:
        vr = vol_cur / vol_prev
        if vr > 1.5: vs += 10
        elif vr > 1.2: vs += 5
    if obv_cur > obv_prev: vs += 5
    if vol_prev > 0 and vol_cur / vol_prev < 0.5: vs -= 5
    score['volume'] = max(0, min(20, vs))
    
    # 波动率分 (0-15)
    vls = 5
    atr = v(last.get('atr'))
    if close > 0 and atr > 0:
        atr_pct = atr / close * 100
        if 1.5 <= atr_pct <= 3.5: vls += 5
        elif atr_pct < 1: vls += 3
        elif atr_pct > 5: vls -= 5
    if len(bars) >= 20:
        cls = [v(b.get('close')) for b in bars[-20:]]
        if min(cls) > 0:
            vol_range = (max(cls) - min(cls)) / min(cls) * 100
            if vol_range < 10: vls += 5
            elif vol_range > 30: vls -= 5
    score['volatility'] = max(0, min(15, vls))
    
    # 风险分 (0-20)
    rs = 10
    if len(bars) >= 10:
        low5 = min(bars[-5:], key=lambda x: v(x.get('low')))
        low10 = min(bars[-10:-5], key=lambda x: v(x.get('low'))) if len(bars) >= 10 else low5
        if v(low5.get('close')) < v(low10.get('close')) and v(low5.get('macd_hist')) > v(low10.get('macd_hist')):
            rs += 5
        if v(low5.get('close')) > v(low10.get('close')) and v(low5.get('macd_hist')) < v(low10.get('macd_hist')):
            rs -= 5
    if ma20 > 0:
        dist = (close - ma20) / ma20 * 100
        if dist < 2: rs += 5
        elif dist > 15: rs -= 5
    score['risk'] = max(0, min(20, rs))
    
    total = sum(score.values())
    return total, score


def compute_kelly(total_score, target_pct, stop_pct):
    """半 Kelly 仓位"""
    p = min(0.6, total_score / 100)
    if stop_pct <= 0:
        return 0
    b = target_pct / stop_pct
    if b <= 1:
        return 0
    kelly = (p * b - (1 - p)) / b
    return max(0, min(0.5, kelly)) * 0.5  # 半 Kelly


def run_strategy_research(start_date, end_date, capital=1000000):
    """主回测入口"""
    conn = sqlite3.connect(DB_PATH)
    
    # 1. 获取所有有历史数据的股票
    stocks = conn.execute("""
        SELECT DISTINCT sd.code, COALESCE(s.name, sd.code) as name
        FROM stock_daily sd
        LEFT JOIN stocks s ON sd.code = s.code
        WHERE sd.date>=? AND sd.date<=?
    """, (start_date, end_date)).fetchall()
    
    results = {
        'period': f"{start_date} ~ {end_date}",
        'capital': capital,
        'total_stocks_screened': len(stocks),
        'scored_stocks': 0,
        'buy_signals': 0,
        'trades': [],
        'summary': {}
    }
    
    trades = []
    code_names = {}
    eval_dates = []  # collect evaluation dates
    
    for code, name in stocks:
        code_names[code] = name
        bars = prepare_bars(code, start_date, end_date)
        if not bars or len(bars) < 20:
            continue
        
        # 2. 滚动评估：每周评分一次
        window_start = 20  # need at least 20 bars for indicators
        while window_start < len(bars):
            window = bars[:window_start]
            score_result = compute_single_score(window)
            if score_result is None:
                window_start += 5
                continue
            total_score, score_components = score_result
            results['scored_stocks'] += 1
            
            last = window[-1]
            close = v_close = last.get('close', 0) or 0
            ma20 = last.get('ma20', 0) or 0
            
            # 3. 买入条件：总评分 >= 45 且 momentum >= 8
            buy_signal = total_score >= 45 and score_components['momentum'] >= 8
            
            if buy_signal:
                results['buy_signals'] += 1
                entry_price = close
                atr_val = last.get('atr', 0) or 0
                if atr_val > 0:
                    target_price = entry_price * 1.10
                    stop_price = entry_price - atr_val * 2.0
                else:
                    target_price = entry_price * 1.10
                    stop_price = entry_price * 0.93
                
                kelly = compute_kelly(total_score, 0.10, (entry_price - stop_price) / entry_price)
                
                # 4. 模拟未来走势
                future = bars[window_start:min(window_start+20, len(bars))]
                exit_price = None
                exit_reason = None
                max_high = entry_price
                min_low = entry_price
                
                for fb in future:
                    fb_high = fb.get('high', 0) or 0
                    fb_low = fb.get('low', 0) or 0
                    fb_close = fb.get('close', 0) or 0
                    max_high = max(max_high, fb_high)
                    min_low = min(min_low, fb_low)
                    
                    if fb_high >= target_price:
                        exit_price = target_price
                        exit_reason = 'target'
                        break
                    elif fb_low <= stop_price:
                        exit_price = fb_close
                        exit_reason = 'stop'
                        break
                
                if exit_price is None:
                    if future:
                        exit_price = future[-1].get('close', entry_price)
                    else:
                        exit_price = entry_price
                    exit_reason = 'keep'
                
                hold_days = 0
                if exit_reason and exit_reason != 'keep' and future:
                    for idx, fb in enumerate(future):
                        fb_high = fb.get('high', 0) or 0
                        fb_low = fb.get('low', 0) or 0
                        if (exit_reason == 'target' and fb_high >= target_price) or \
                           (exit_reason == 'stop' and fb_low <= stop_price):
                            hold_days = idx + 1
                            break
                elif exit_reason == 'keep':
                    hold_days = len(future)
                
                profit_pct = (exit_price - entry_price) / entry_price * 100 if entry_price > 0 else 0
                
                trades.append({
                    'code': code,
                    'name': name,
                    'entry_date': last.get('date'),
                    'entry_price': round(entry_price, 2),
                    'exit_price': round(exit_price, 2),
                    'profit_pct': round(profit_pct, 2),
                    'exit_reason': exit_reason,
                    'hold_days': hold_days,
                    'score': total_score,
                    'kelly': round(kelly, 3),
                    'stop_loss': round(stop_price, 2),
                    'target': round(target_price, 2),
                    'max_high': round(max_high, 2),
                    'min_low': round(min_low, 2)
                })
            
            window_start += 5  # 每 5 个交易日评估一次
    
    conn.close()
    
    # 5. 汇总
    if trades:
        profits = [t['profit_pct'] for t in trades]
        wins = [t for t in trades if t['profit_pct'] > 0]
        losses = [t for t in trades if t['profit_pct'] <= 0]
        
        win_rate = len(wins) / len(trades) * 100 if trades else 0
        avg_profit = sum(profits) / len(profits) if profits else 0
        avg_win = sum(t['profit_pct'] for t in wins) / len(wins) if wins else 0
        avg_loss = sum(t['profit_pct'] for t in losses) / len(losses) if losses else 0
        
        mean_ret = avg_profit / 100
        if len(profits) > 1:
            variance = sum((p/100 - mean_ret)**2 for p in profits) / (len(profits) - 1)
            std_ret = math.sqrt(variance)
        else:
            std_ret = 0
        sharpe = mean_ret / std_ret * math.sqrt(252) if std_ret > 0 else 0
        
        # 最大回撤：用实际资本曲线计算
        capital_curve = [capital]
        for t in trades:
            new_cap = capital_curve[-1] * (1 + t['profit_pct'] / 100)
            capital_curve.append(new_cap)
        peak_cap = capital
        max_dd = 0
        for c in capital_curve:
            peak_cap = max(peak_cap, c)
            dd = (peak_cap - c) / peak_cap * 100
            max_dd = max(max_dd, dd)
        
        results['trades'] = sorted(trades, key=lambda x: abs(x['profit_pct']), reverse=True)[:50]
        results['summary'] = {
            'total_trades': len(trades),
            'win_rate': round(win_rate, 1),
            'avg_profit_pct': round(avg_profit, 2),
            'avg_win_pct': round(avg_win, 2),
            'avg_loss_pct': round(avg_loss, 2),
            'sharpe_ratio': round(sharpe, 2),
            'max_drawdown_pct': round(max_dd, 2),
            'profit_factor': round(abs(avg_win / avg_loss), 2) if avg_loss != 0 else float('inf'),
            'wins': len(wins),
            'losses': len(losses),
            'capital_end': round(capital * (1 + avg_profit/100), 2)
        }
    
    return results


if __name__ == '__main__':
    end = datetime.now().strftime('%Y-%m-%d')
    start = (datetime.now() - timedelta(days=180)).strftime('%Y-%m-%d')
    r = run_strategy_research(start, end)
    print(json.dumps(r, indent=2, ensure_ascii=False))