#!/usr/bin/env python3
"""
小唯股票策略对比 — 趋势跟踪 vs MACD
====================================
均线突破策略(MA_Breakout): 收盘价上穿20日均线买入，下穿卖出

用法:
  python3 stock_compare.py <股票代码> [起始] [结束]
"""

import json, sys, urllib.request
from datetime import datetime
from pathlib import Path

import numpy as np
import pandas as pd

OUTPUT = Path.home() / ".hermes" / "stock_backtest"
OUTPUT.mkdir(exist_ok=True)
FEISHU_WEBHOOK = "https://open.feishu.cn/open-apis/bot/v2/hook/446db983-e392-4d2c-bfb8-f9060e5df3ad"


def send_feishu(msg):
    payload = json.dumps({"msg_type": "text", "content": {"text": msg}}).encode()
    req = urllib.request.Request(FEISHU_WEBHOOK, data=payload,
        headers={"Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=10):
            pass
    except Exception:
        pass


def get_data(code, start, end):
    mc = f"sh{code}" if code.startswith("6") else f"sz{code}"
    url = (f"https://web.ifzq.gtimg.cn/appstock/app/fqkline/get"
           f"?_var=kline_dayqfq&param={mc},day,{start},{end},500,qfq")
    try:
        text = urllib.request.urlopen(url, timeout=10).read().decode("utf-8")
        data = json.loads(text.replace("kline_dayqfq=", "", 1))
        qfq = (data.get("data", {}).get(mc, {}).get("qfqday") or
               data.get("data", {}).get(mc, {}).get("day") or [])
        rows = []
        for item in qfq:
            if len(item) < 6:
                continue
            try:
                rows.append({"date": item[0], "open": float(item[1]),
                    "close": float(item[2]), "high": float(item[3]),
                    "low": float(item[4]), "volume": float(item[5])})
            except (ValueError, IndexError):
                continue
        df = pd.DataFrame(rows)
        df["date"] = pd.to_datetime(df["date"])
        df.set_index("date", inplace=True)
        df.sort_index(inplace=True)
        return df
    except Exception:
        return None


def macd_backtest(df, fast=12, slow=26, sig=9, cash=100000):
    ema_f = df["close"].ewm(span=fast).mean()
    ema_s = df["close"].ewm(span=slow).mean()
    macd = ema_f - ema_s
    signal = macd.ewm(span=sig).mean()

    shares = 0; c = cash; peak = cash; max_dd = 0
    trades = []
    for i in range(slow, len(df)):
        p = df["close"].iloc[i]
        if macd.iloc[i] > signal.iloc[i] and macd.iloc[i-1] <= signal.iloc[i-1]:
            if shares == 0:
                n = int(c / p); c -= n * p; shares = n
                trades.append(("BUY", df.index[i], n, p))
        elif macd.iloc[i] < signal.iloc[i] and macd.iloc[i-1] >= signal.iloc[i-1]:
            if shares > 0:
                c += shares * p; trades.append(("SELL", df.index[i], shares, p)); shares = 0
        peak = max(peak, c + shares * p)
        dd = (peak - (c + shares * p)) / peak * 100 if peak > 0 else 0
        max_dd = max(max_dd, dd)
    final = c + shares * df["close"].iloc[-1]
    wins = len([t for t in trades if t[0] == "SELL" and t[3] > 0])
    return dict(
        final=final, ret=(final-cash)/cash*100,
        buyhold=(df["close"].iloc[-1]-df["close"].iloc[0])/df["close"].iloc[0]*100,
        max_dd=max_dd, trades=len(trades)//2,
        winrate=wins/(len(trades)//2)*100 if trades else 0
    )


def ma_breakout_backtest(df, ma_days=20, cash=100000):
    ma = df["close"].rolling(ma_days).mean()
    shares = 0; c = cash; peak = cash; max_dd = 0
    trades = []
    for i in range(ma_days, len(df)):
        p = df["close"].iloc[i]
        if df["close"].iloc[i] > ma.iloc[i] and df["close"].iloc[i-1] <= ma.iloc[i-1]:
            if shares == 0:
                n = int(c / p); c -= n * p; shares = n
                trades.append(("BUY", df.index[i], n, p))
        elif df["close"].iloc[i] < ma.iloc[i] and df["close"].iloc[i-1] >= ma.iloc[i-1]:
            if shares > 0:
                c += shares * p; trades.append(("SELL", df.index[i], shares, p)); shares = 0
        peak = max(peak, c + shares * p)
        dd = (peak - (c + shares * p)) / peak * 100 if peak > 0 else 0
        max_dd = max(max_dd, dd)
    final = c + shares * df["close"].iloc[-1]
    wins = len([t for t in trades if t[0] == "SELL" and t[3] > 0])
    return dict(
        final=final, ret=(final-cash)/cash*100,
        buyhold=(df["close"].iloc[-1]-df["close"].iloc[0])/df["close"].iloc[0]*100,
        max_dd=max_dd, trades=len(trades)//2,
        winrate=wins/(len(trades)//2)*100 if trades else 0
    )


def dual_ma_backtest(df, fast=5, slow=20, cash=100000):
    """双均线策略: 快线穿慢线金叉买，死叉卖"""
    ma_fast = df["close"].rolling(fast).mean()
    ma_slow = df["close"].rolling(slow).mean()
    shares = 0; c = cash; peak = cash; max_dd = 0
    trades = []
    for i in range(slow, len(df)):
        p = df["close"].iloc[i]
        if ma_fast.iloc[i] > ma_slow.iloc[i] and ma_fast.iloc[i-1] <= ma_slow.iloc[i-1]:
            if shares == 0:
                n = int(c / p); c -= n * p; shares = n
                trades.append(("BUY", df.index[i], n, p))
        elif ma_fast.iloc[i] < ma_slow.iloc[i] and ma_fast.iloc[i-1] >= ma_slow.iloc[i-1]:
            if shares > 0:
                c += shares * p; trades.append(("SELL", df.index[i], shares, p)); shares = 0
        peak = max(peak, c + shares * p)
        dd = (peak - (c + shares * p)) / peak * 100 if peak > 0 else 0
        max_dd = max(max_dd, dd)
    final = c + shares * df["close"].iloc[-1]
    wins = len([t for t in trades if t[0] == "SELL" and t[3] > 0])
    return dict(
        final=final, ret=(final-cash)/cash*100,
        buyhold=(df["close"].iloc[-1]-df["close"].iloc[0])/df["close"].iloc[0]*100,
        max_dd=max_dd, trades=len(trades)//2,
        winrate=wins/(len(trades)//2)*100 if trades else 0
    )


def main():
    if len(sys.argv) < 2:
        print("用法: python3 stock_compare.py <股票代码> [起始] [结束]")
        sys.exit(1)

    code = sys.argv[1]
    start = sys.argv[2] if len(sys.argv) > 2 else "2023-01-01"
    end = sys.argv[3] if len(sys.argv) > 3 else datetime.now().strftime("%Y-%m-%d")

    print(f"\n代码: {code} | {start} ~ {end}")
    df = get_data(code, start, end)
    if df is None or len(df) < 60:
        print("数据获取失败"); sys.exit(1)
    print(f"数据: {len(df)}条")

    r1 = macd_backtest(df)
    r2 = ma_breakout_backtest(df)
    r3 = dual_ma_backtest(df)

    print(f"\n{'='*60}")
    print(f"{'策略':<20} {'收益':>10} {'买入持有':>10} {'α':>10} {'最大回撤':>10} {'交易':>6} {'胜率':>8}")
    print(f"{'-'*60}")
    print(f"{'MACD(12,26,9)':<20} {r1['ret']:>+9.1f}% {r1['buyhold']:>+9.1f}% {r1['ret']-r1['buyhold']:>+9.1f}% {r1['max_dd']:>9.1f}% {r1['trades']:>6} {r1['winrate']:>7.0f}%")
    print(f"{'MA突破(20日)':<20} {r2['ret']:>+9.1f}% {r2['buyhold']:>+9.1f}% {r2['ret']-r2['buyhold']:>+9.1f}% {r2['max_dd']:>9.1f}% {r2['trades']:>6} {r2['winrate']:>7.0f}%")
    print(f"{'双均线(5,20)':<20} {r3['ret']:>+9.1f}% {r3['buyhold']:>+9.1f}% {r3['ret']-r3['buyhold']:>+9.1f}% {r3['max_dd']:>9.1f}% {r3['trades']:>6} {r3['winrate']:>7.0f}%")
    print(f"{'='*60}")

    # 飞书
    msg = f"""📊 策略对比报告

代码: {code} | {start} ~ {end} | {len(df)}条

策略收益对比:
MACD(12,26,9):    {r1['ret']:+.1f}% (α={r1['ret']-r1['buyhold']:+.1f}%, 回撤{r1['max_dd']:.1f}%, {r1['trades']}笔, 胜率{r1['winrate']:.0f}%)
MA突破(20日):    {r2['ret']:+.1f}% (α={r2['ret']-r2['buyhold']:+.1f}%, 回撤{r2['max_dd']:.1f}%, {r2['trades']}笔, 胜率{r2['winrate']:.0f}%)
双均线(5,20):    {r3['ret']:+.1f}% (α={r3['ret']-r3['buyhold']:+.1f}%, 回撤{r3['max_dd']:.1f}%, {r3['trades']}笔, 胜率{r3['winrate']:.0f}%)

买入持有基准: {r1['buyhold']:+.1f}%

生成: {datetime.now().strftime('%Y-%m-%d %H:%M')}
小唯股票投研 · 模拟阶段"""
    send_feishu(msg)
    print("\n✅ 已推送飞书")

    # 保存
    result_file = OUTPUT / f"compare_{code}.json"
    with open(result_file, "w") as f:
        json.dump({"code": code, "start": start, "end": end,
            "macd": r1, "ma_breakout": r2, "dual_ma": r3}, f, ensure_ascii=False, indent=2, default=str)
    print(f"数据存: {result_file}")


if __name__ == "__main__":
    main()